[cherry-pick] fix propagation of mask creator objects in PT pruning modifier (#220) (#221)

* fix propagation of mask creator objects in PT pruning modifier (#220)

* fix propagation of mask creator objects in PT pruning modifier

* quality - need type() to handle inheretence

* update version to 0.3.1 for hotfix

Author: bfineran

src/sparseml/pytorch/optim/modifier_pruning.py

@@ -608,6 +608,7 @@ def initialize(self, module: Module, optimizer: Optimizer):
             layers,
             param_names,
             layer_names=layer_names,
+            mask_creator=self._mask_creator,
             global_sparsity=self._global_sparsity,
         )
 

tests/sparseml/pytorch/optim/test_modifier_pruning.py

@@ -23,6 +23,7 @@
     GlobalMagnitudePruningModifier,
     GMPruningModifier,
     MagnitudePruningModifier,
+    load_mask_creator,
 )
 from tests.sparseml.pytorch.helpers import LinearNet
 from tests.sparseml.pytorch.optim.test_modifier import (
@@ -262,6 +263,10 @@ def test_lifecycle(
         optimizer = optim_lambda(model)
         self.initialize_helper(modifier, model, optimizer)
         assert modifier.applied_sparsity is None
+        assert type(load_mask_creator(modifier._mask_type)) == type(  # noqa: E721
+            modifier._mask_creator
+        )
+        assert modifier._mask_creator == modifier._module_masks._mask_creator
 
         # check sparsity is not set before
         for epoch in range(int(modifier.start_epoch)):