fix test_process_mesh_v2 ci fail!

PaddlePaddle · Feb 29, 2024 · 82fd7ae · 82fd7ae
1 parent aa33c8a
commit 82fd7ae
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h b/paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h
@@ -52,7 +52,9 @@ struct ProcessMeshAttrStorage : public pir::AttributeStorage {
   ///
   /// \brief Each derived TypeStorage needs to overload operator==.
   ///
-  bool operator==(const ParamKey& key) const { return process_mesh == key; }
+  bool operator==(const ParamKey& key) const {
+    return process_mesh == key && process_mesh.dim_names() == key.dim_names();
+  }
 
   ParamKey process_mesh;
 };

diff --git a/paddle/phi/core/distributed/auto_parallel/process_mesh.cc b/paddle/phi/core/distributed/auto_parallel/process_mesh.cc
@@ -127,9 +127,6 @@ bool operator==(const ProcessMesh &lhs, const ProcessMesh &rhs) {
   if (lhs.process_ids() != rhs.process_ids()) {
     return false;
   }
-  if (lhs.dim_names() != rhs.dim_names()) {
-    return false;
-  }
   return true;
 }