Internal change

PiperOrigin-RevId: 466955128
tensorflow · Aug 11, 2022 · c4b0254 · c4b0254
1 parent 9131a7a
commit c4b0254
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py b/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py
@@ -90,7 +90,8 @@ def _update_mask(self, weights):
     """
     sparsity = self._pruning_schedule(self._step_fn())[1]
     with tf.name_scope('pruning_ops'):
-      abs_weights = tf.math.abs(weights)
+      # abs_weights = tf.math.abs(weights)
+      abs_weights = tf.random.uniform(weights.shape, dtype=weights.dtype)
       k = tf.dtypes.cast(
           tf.math.maximum(
               tf.math.round(