fix reshard dist_attr

Author: LiYuRio

paddle/phi/core/distributed/auto_parallel/dist_tensor.cc

@@ -117,6 +117,9 @@ DistTensor::DistTensor() : value_(std::make_shared<DenseTensor>()) {}
 DistTensor::DistTensor(const std::shared_ptr<phi::DenseTensor>& global_value,
                        const TensorDistAttr& dist_attr)
     : global_dims_(global_value->dims()), dist_attr_(dist_attr) {
+  process_mesh_ = dist_attr_.process_mesh();
+  placements_ = ToPlacements(dist_attr);
+
   // If the current rank doesn't in process_mesh, we should create an
   // uninitialized tensor only with tensor_meta.
   if (IsCurRankInMesh(dist_attr.process_mesh())) {

paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.cc

@@ -52,6 +52,8 @@ void ReshardFunction::SetDistProps(DistTensor* tensor,
 
   tensor->global_dims_ = dims;
   tensor->dist_attr_ = dist_attr;
+  tensor->process_mesh_ = dist_attr.process_mesh();
+  tensor->placements_ = ToPlacements(dist_attr);
 }
 
 void ReshardFunction::SetDistProps(DistTensor* tensor,
@@ -64,6 +66,8 @@ void ReshardFunction::SetDistProps(DistTensor* tensor,
                         str_join(vectorize(tensor->dims()))));
 
   tensor->dist_attr_ = dist_attr;
+  tensor->process_mesh_ = dist_attr.process_mesh();
+  tensor->placements_ = ToPlacements(dist_attr);
 }
 
 DenseTensor* ReshardFunction::GetMutableTensor(DistTensor* tensor) {