index select support auto parallel (#58841)

PaddlePaddle · Nov 10, 2023 · f2f8879 · f2f8879
1 parent 2a45bea
commit f2f8879
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -1442,6 +1442,10 @@ static PyObject* tensor__getitem_index_not_tensor(TensorObject* self,
           egr::Controller::Instance().GetExpectedPlace());
       paddle::framework::TensorFromVector(
           list_select_idxs, *dev_ctx, idx_tensor.get());
+      const phi::distributed::ProcessMesh* mesh = nullptr;
+      if (InputsContainDistTensor(&mesh, self->tensor, select_index)) {
+        ConvertAllInputsToDistTensor(mesh, self->tensor, select_index);
+      }
       out = index_select_ad_func(self->tensor, select_index, 0);
     }
   }