Fix skip_if_not_multigpu decorator (#54916)

Summary: Pull Request resolved: #54916 Fixes #54887 `skip_if_not_multigpu` was skipping all the tests that use it. Test Plan: Imported from OSS Reviewed By: mrshenli Differential Revision: D27412193 Pulled By: H-Huang fbshipit-source-id: 28d6697bd8cc6b6784cdb038ccb3ff138d0610eb
pytorch · Apr 2, 2021 · 5610e82 · 5610e82
1 parent 8822c7e
commit 5610e82
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 75 deletions.
diff --git a/test/distributed/optim/test_zero_redundancy_optimizer.py b/test/distributed/optim/test_zero_redundancy_optimizer.py
@@ -347,7 +347,7 @@ def some_trainable():
         all_trainable()
         some_trainable()
 
-    @common_distributed.skip_if_not_multigpu
+    @common_distributed.skip_if_lt_x_gpu(2)
     def test_collect_shards(self):
         """ Check the state consolidation mechanism, and the state dict exposed by ZeroRedundancyOptimizer"""
         self.dist_init(self.rank)