test cpu

pytorch · Nov 26, 2023 · 8f3fdfc · 8f3fdfc
1 parent 9666e5f
commit 8f3fdfc
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/test/spmd/test_spmd_debugging.py b/test/spmd/test_spmd_debugging.py
@@ -30,8 +30,8 @@ def setUpClass(cls):
 
   @unittest.skipIf(
       not xr.using_pjrt() or
-      xu.getenv_as(xenv.PJRT_DEVICE, str) in ("GPU", 'CUDA', 'ROCM', 'CPU'),
-      f"Requires PJRT_DEVICE set to `TPU`.")
+      xu.getenv_as(xenv.PJRT_DEVICE, str) in ("GPU", 'CUDA', 'ROCM'),
+      f"Requires PJRT_DEVICE set to `TPU`, `CPU`.")
   def test_debugging_spmd_single_host_tiled(self):
     from torch_xla.distributed.spmd.debugging import visualize_tensor_sharding
     device = xm.xla_device()
@@ -108,8 +108,8 @@ def test_debugging_spmd_single_host_tiled(self):
 
   @unittest.skipIf(
       not xr.using_pjrt() or
-      xu.getenv_as(xenv.PJRT_DEVICE, str) in ("GPU", 'CUDA', 'ROCM', 'CPU'),
-      f"Requires PJRT_DEVICE set to `TPU`.")
+      xu.getenv_as(xenv.PJRT_DEVICE, str) in ("GPU", 'CUDA', 'ROCM'),
+      f"Requires PJRT_DEVICE set to `TPU`, 'CPU'.")
   def test_single_host_partial_replication(self):
     from torch_xla.distributed.spmd.debugging import visualize_tensor_sharding
     device = xm.xla_device()
@@ -158,8 +158,8 @@ def test_single_host_partial_replication(self):
 
   @unittest.skipIf(
       not xr.using_pjrt() or
-      xu.getenv_as(xenv.PJRT_DEVICE, str) in ("GPU", 'CUDA', 'ROCM', 'CPU'),
-      f"Requires PJRT_DEVICE set to `TPU`.")
+      xu.getenv_as(xenv.PJRT_DEVICE, str) in ("GPU", 'CUDA', 'ROCM'),
+      f"Requires PJRT_DEVICE set to `TPU`, 'CPU'.")
   def test_single_host_replicated(self):
     from torch_xla.distributed.spmd.debugging import visualize_tensor_sharding
     device = xm.xla_device()