diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc index 3877200ad310e..78a260f35fb84 100644 --- a/paddle/fluid/pybind/eager_method.cc +++ b/paddle/fluid/pybind/eager_method.cc @@ -675,7 +675,12 @@ static PyObject* tensor_method_copy_(TensorObject* self, PyObject* args, PyObject* kwargs) { EAGER_TRY - paddle::Tensor src_tensor = CastPyArg2Tensor(PyTuple_GET_ITEM(args, 0), 0); + paddle::Tensor& src_tensor = CastPyArg2Tensor(PyTuple_GET_ITEM(args, 0), 0); + const phi::distributed::ProcessMesh* mesh = nullptr; + if (InputsContainDistTensor(&mesh, src_tensor, self->tensor)) { + ConvertAllInputsToDistTensor(mesh, src_tensor, self->tensor); + } + bool blocking = CastPyArg2AttrBoolean(PyTuple_GET_ITEM(args, 1), 1); VLOG(6) << "Start Copy Tensor " << src_tensor.name() << " to " << self->tensor.name(); diff --git a/paddle/fluid/pybind/eager_properties.cc b/paddle/fluid/pybind/eager_properties.cc index a9d7b03556cfa..582d15909e941 100644 --- a/paddle/fluid/pybind/eager_properties.cc +++ b/paddle/fluid/pybind/eager_properties.cc @@ -300,7 +300,7 @@ int tensor_properties_set_grad(TensorObject* self, PyObject* value, void* closure) { EAGER_TRY - auto src = CastPyArg2Tensor(value, 0); + auto& src = CastPyArg2Tensor(value, 0); PADDLE_ENFORCE( egr::EagerUtils::IsLeafTensor(self->tensor), paddle::platform::errors::Fatal("Only leaf Tensor can be set grad.")); @@ -311,6 +311,10 @@ int tensor_properties_set_grad(TensorObject* self, "Detected NULL grad" "Please check if you have manually cleared" "the grad inside autograd_meta")); + const phi::distributed::ProcessMesh* mesh = nullptr; + if (InputsContainDistTensor(&mesh, src, self->tensor, *grad)) { + ConvertAllInputsToDistTensor(mesh, src, self->tensor, *grad); + } grad->copy_(src, self->tensor.place(), true); return 0; EAGER_CATCH_AND_THROW_RETURN_NEG