Merge pull request #860 from psunn:psunn_quickfix_mbyn

PiperOrigin-RevId: 401647817

Author: tensorflower-gardener

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

@@ -240,7 +240,11 @@ def conditional_mask_update(self):
     """Returns an op to updates masks as per the pruning schedule."""
 
     def maybe_update_masks():
-      return self._pruning_schedule(self._step_fn())[0]
+      if self._sparsity_m_by_n:
+        # Update structured sparsity masks only at step 1
+        return tf.math.equal(self._step_fn(), 1)
+      else:
+        return self._pruning_schedule(self._step_fn())[0]
 
     def no_update():
       return tf.no_op()