Update

kwen2501 · kwen2501 · commit 95f6d5dd8cda · 2025-05-23T23:01:28.000-07:00
[ghstack-poisoned]
diff --git a/test/distributed/pipelining/test_stage.py b/test/distributed/pipelining/test_stage.py
@@ -24,6 +24,7 @@
     instantiate_parametrized_tests,
     parametrize,
     run_tests,
+    skip_but_pass_in_sandcastle,
     skip_but_pass_in_sandcastle_if,
 )
 from torch.utils._pytree import tree_map_only
@@ -68,6 +69,10 @@ def backend_str(cls) -> str:
         # Testing with NCCL backend
         return "nccl"
 
+    @classmethod
+    def device_type(cls) -> str:
+        return device_type
+
     @property
     def device(self) -> torch.device:
         return torch.device(device_type, self.rank)
@@ -350,7 +355,7 @@ def init_pg(self):
         )
 
     @requires_nccl()
-    @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
+    @skip_but_pass_in_sandcastle("Flaky in CI")
     def test_shape_prop_mismatch(self):
         """Tests shape prop errors are raised"""
         self.init_pg()
diff --git a/torch/testing/_internal/common_distributed.py b/torch/testing/_internal/common_distributed.py
@@ -1529,7 +1529,10 @@ def backend_str(cls) -> Optional[str]:
     # Please override if you intend to test on specific device type
     @classmethod
     def device_type(cls) -> str:
-        return torch.accelerator.current_accelerator().type
+        curr_device = torch.accelerator.current_accelerator()
+        if curr_device is None:
+            return "cpu"
+        return curr_device.type
 
     @classmethod
     def opts(cls, high_priority_stream=False):
