From 066097e83f564c657f91ad2ed64a80098973fe44 Mon Sep 17 00:00:00 2001 From: ronnywang Date: Mon, 14 Aug 2023 13:30:47 +0800 Subject: [PATCH] Wait the stream before memcpy_d2h_sync in tensor.numpy (#56189) --- paddle/fluid/pybind/eager_method.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc index e78e3bb106d..6f49953ae8a 100644 --- a/paddle/fluid/pybind/eager_method.cc +++ b/paddle/fluid/pybind/eager_method.cc @@ -272,6 +272,7 @@ static PyObject* tensor_method_numpy(TensorObject* self, gpuMemcpyKind kind = cudaMemcpyDeviceToHost; #elif defined(PADDLE_WITH_HIP) gpuMemcpyKind kind = hipMemcpyDeviceToHost; + phi::DeviceContextPool::Instance().Get(self->tensor.place())->Wait(); #endif if (self->tensor.is_selected_rows()) { VLOG(6) << "Getting SelectedRows's numpy value"; @@ -341,6 +342,7 @@ static PyObject* tensor_method_numpy(TensorObject* self, #ifdef PADDLE_WITH_CUSTOM_DEVICE } else if (self->tensor.is_custom_device()) { eager_gil_scoped_release guard; + phi::DeviceContextPool::Instance().Get(self->tensor.place())->Wait(); if (self->tensor.is_selected_rows()) { VLOG(6) << "Getting SelectedRows's numpy value"; auto* selected_rows = -- GitLab