[AutoParallel] fix process_mesh (PaddlePaddle#46583)

zhaoyinglia · Oct 19, 2022 · 13a5bca · 13a5bca
1 parent c2b47e4
commit 13a5bca
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/python/paddle/distributed/auto_parallel/process_mesh.py b/python/paddle/distributed/auto_parallel/process_mesh.py
@@ -168,7 +168,10 @@ def __getitem__(self, index):
         else:
             new_mesh = self._mesh[index]
             new_dim_names = self._dim_names[1:]
-            return ProcessMesh(new_mesh, new_dim_names)
+            if new_mesh.shape:
+                return ProcessMesh(new_mesh, new_dim_names)
+            else:
+                return ProcessMesh([new_mesh])
 
     def __enter__(self):
         set_current_process_mesh(self)

diff --git a/python/paddle/fluid/tests/unittests/auto_parallel/test_process_mesh.py b/python/paddle/fluid/tests/unittests/auto_parallel/test_process_mesh.py
@@ -101,6 +101,12 @@ def test_construction(self):
         self.assertEqual(sub_process_mesh4.dim_names, ["d0"])
         self.assertEqual(sub_process_mesh4.ndim, 1)
 
+        sub_process_mesh5 = sub_process_mesh3[0]
+        self.assertEqual(sub_process_mesh5.shape, [1])
+        self.assertEqual(sub_process_mesh5.process_ids, [1])
+        self.assertEqual(sub_process_mesh5.dim_names, ["d0"])
+        self.assertEqual(sub_process_mesh5.ndim, 1)
+
     def test_context_manager(self):
         mesh = np.array([1, 2, 3, 4])
         input = static.data(name="input",