const pruning mask was inverted (#1751)

neuralmagic · Oct 6, 2023 · 528289b · 528289b
1 parent 99a60ab
commit 528289b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/sparseml/modifiers/pruning/constant/pytorch.py b/src/sparseml/modifiers/pruning/constant/pytorch.py
@@ -54,7 +54,7 @@ def on_finalize(self, state: State, **kwargs) -> bool:
     def on_start(self, state: State, event: Event, **kwargs):
         for layer_param_name, parameterized_layer in self.parameterized_layers_.items():
             self.update_mask(
-                layer_param_name, parameterized_layer.param.data.abs() < self._epsilon
+                layer_param_name, parameterized_layer.param.data.abs() > self._epsilon
             )
 
         self.enable_masks()