From c4b0254469b5563ed56d54b37f1c6c51e0e0d947 Mon Sep 17 00:00:00 2001 From: Utku Evci Date: Thu, 11 Aug 2022 08:06:05 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 466955128 --- .../python/core/sparsity/keras/pruning_impl.py | 3 ++- 1 file changed, 2 insertions(+), 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 index f836f36ac..a097d4ede 100644 --- 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(