[TestZeroRedundancyOptimizer] Add multi gpu checker

jaglinux · jaglinux · commit 4b0df0b6bf3d · 2021-03-08T23:12:55.000Z
The test test_collect_shards fails on single GPU setup.
Enabling the multi gpu checker.

Signed-off-by: Jagadish Krishnamoorthy &lt;jagdish.krishna@gmail.com&gt;
diff --git a/test/distributed/optim/test_zero_redundancy_optimizer.py b/test/distributed/optim/test_zero_redundancy_optimizer.py
@@ -346,6 +346,7 @@ def some_trainable():
         all_trainable()
         some_trainable()
 
+    @common_distributed.skip_if_not_multigpu
     def test_collect_shards(self):
         """ Check the state consolidation mechanism, and the state dict exposed by ZeroRedundancyOptimizer"""
         self.dist_init(self.rank)
