Change pruning_step type to int64 from int32.

PiperOrigin-RevId: 278802043

Author: Xhark

tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py

@@ -201,7 +201,7 @@ def build(self, input_shape):
         'pruning_step',
         shape=[],
         initializer=initializers.Constant(-1),
-        dtype=dtypes.int32,
+        dtype=dtypes.int64,
         trainable=False)
 
     def training_step_fn():
@@ -222,7 +222,9 @@ def call(self, inputs, training=None):
     def add_update():
       with ops.control_dependencies([
           check_ops.assert_greater_equal(
-              self.pruning_step, 0, message=self._PRUNE_CALLBACK_ERROR_MSG)]):
+              self.pruning_step,
+              np.int64(0),
+              message=self._PRUNE_CALLBACK_ERROR_MSG)]):
         with ops.control_dependencies(
             [self.pruning_obj.conditional_mask_update()]):
           return control_flow_ops.no_op('update')