[sparsity] sparsifier step tests (pytorch#60107)

Summary:
Pull Request resolved: pytorch#60107

Unit tests for sparsifier `step`

Test Plan:
`buck test mode/dev-nosan //caffe2/test:ao -- TestWeightNormSparsifier`

https://pxl.cl/1LhQP

Reviewed By: z-a-f

Differential Revision: D29167029

fbshipit-source-id: 053027ca92701097406372ef0f81d79ef28380aa

test/ao/sparsity/test_sparsifier.py

@@ -152,9 +152,17 @@ def test_step(self):
         model = Model()
         sparsifier = WeightNormSparsifier(sparsity_level=0.5)
         sparsifier.prepare(model, config=[model.linear])
+        for g in sparsifier.module_groups:
+            # Before step
+            module = g['module']
+            assert (1.0 - module.parametrizations['weight'][0].mask.mean()) == 0  # checking sparsity level is 0
         sparsifier.enable_mask_update = True
         sparsifier.step()
         self.assertAlmostEqual(model.linear.parametrizations['weight'][0].mask.mean().item(), 0.5, places=2)
+        for g in sparsifier.module_groups:
+            # After step
+            module = g['module']
+            assert (1.0 - module.parametrizations['weight'][0].mask.mean()) > 0  # checking sparsity level has increased
 
     def test_prepare(self):
         model = Model()