Minor update to the unit test for the upcoming Keras RNG change.

PiperOrigin-RevId: 413582073

Author: qlzh727

tensorflow_model_optimization/python/core/clustering/keras/cluster_integration_test.py

@@ -335,6 +335,10 @@ def on_train_batch_begin(self, batch, logs=None):
         self.bias = (self.model.layers[0].layer.bias.numpy())
 
       def on_train_batch_end(self, batch, logs=None):
+        # Skip the first batch since the update from the start of training is
+        # too small.
+        if not batch:
+          return
         # Check weights are different after batch
         assert not np.array_equal(
             self.original_weight_kernel,