fix BasicShardingTest.test_2d_tensor_3d_mesh on cpu

pytorch · Feb 2, 2024 · 960d609 · 960d609
1 parent 43ef193
commit 960d609
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/test/spmd/test_xla_sharding.py b/test/spmd/test_xla_sharding.py
@@ -728,10 +728,14 @@ def test_2d_tensor_3d_mesh(self):
       mesh = self._get_mesh((2, self.n_devices // 2, 1))
       xs.mark_sharding(t1, mesh, partition_spec=(2, 1))
       sharding_annotation = 'sharding={devices=[1,%d,2]' % (self.n_devices // 2)
-    else:
+    elif self.n_devices == 2:
       mesh = self._get_mesh((2, 1, 1))
       xs.mark_sharding(t1, mesh, partition_spec=(2, 1))
       sharding_annotation = "sharding={replicated}"
+    else:
+      mesh = self._get_mesh((1, 1, 1))
+      xs.mark_sharding(t1, mesh, partition_spec=(2, 1))
+      sharding_annotation = "sharding={replicated}"
     self.assertIn(sharding_annotation,
                   torch_xla._XLAC._get_xla_tensors_hlo([t1]))
     actual = (t1 + t2).cpu()