Strange behaviour pruning with small datasets

Describe the bug
Prune Low Magnitude seems not to update the weights to 0 (I am using Constant Sparsity), when using a small dataset for training (1000 Images).

System information

TensorFlow version: 2.8.1

TensorFlow Model Optimization version (installed from source or binary): 0.7.2

Python version: 3.8.10

Describe the expected behavior
Pruning should be working as normal on small dataset as it is working on bigger datasets.

Describe the current behavior
The weights are not updated to 0 after a model_for_pruning.fit run (see in Code Example). The exact same example works if you increase the size of the dataset (var dataset_size) to 10000 or change the batch_size to 16. I don't think, that this is intended when using the Constant Sparsity Feature or am I doing something wrong here?
Code to reproduce the issue

# Random "images" for test prupose
dataset_size = 1000
train_images = np.random.rand(dataset_size, 128,128,3) 
train_labels = np.random.rand(dataset_size, 10)

model = keras.Sequential([
  keras.layers.InputLayer(input_shape=(128, 128,3)),
  keras.layers.Reshape(target_shape=(128, 128, 3)),
  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),
  keras.layers.MaxPooling2D(pool_size=(2, 2)),
  keras.layers.Flatten(),
  keras.layers.Dense(10)
])
prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

batch_size = 64
epochs = 2
validation_split = 0.1

pruning_params = {
      'pruning_schedule': tfmot.sparsity.keras.ConstantSparsity(0.8, 0)
}

model_for_pruning = prune_low_magnitude(model, **pruning_params)
model_for_pruning.compile(optimizer='adam',
              loss=tf.keras.losses.BinaryCrossentropy(from_logits=False),
              metrics=['accuracy'])

callbacks = [
  tfmot.sparsity.keras.UpdatePruningStep(),
  tfmot.sparsity.keras.PruningSummaries(log_dir=logdir),
]
  
model_for_pruning.fit(train_images, train_labels,
                  batch_size=batch_size, epochs=epochs, validation_split=validation_split,
                  callbacks=callbacks)

print(model_for_pruning.get_weights()) 
#Output contains no pruned weights / zero values