switch gloo tests back to cpu (pytorch#1665)

Summary: Also minor refactoring. Differential Revision: D53149479
henrylhtsang · Jan 30, 2024 · 13e0ae5 · 13e0ae5
1 parent 0e153c9
commit 13e0ae5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrec/distributed/tests/test_model_parallel_gloo.py b/torchrec/distributed/tests/test_model_parallel_gloo.py
@@ -16,9 +16,9 @@ class ModelParallelTestGloo(ModelParallelBase):
     pass
 
 
-class ModelParallelStateDictTestNccl(ModelParallelStateDictBase):
+class ModelParallelStateDictTestGloo(ModelParallelStateDictBase):
     pass
 
 
-class ModelParallelSparseOnlyTestNccl(ModelParallelSparseOnlyBase):
+class ModelParallelSparseOnlyTestGloo(ModelParallelSparseOnlyBase):
     pass