From a822f8ddbbf8e48d133ed78048dc4c17d863af74 Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Mon, 16 Apr 2018 16:23:22 +0800
Subject: [PATCH] Add wait

---
 paddle/fluid/pybind/tensor_py.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/paddle/fluid/pybind/tensor_py.h b/paddle/fluid/pybind/tensor_py.h
index 4a9dbd324c..c9cad15a74 100644
--- a/paddle/fluid/pybind/tensor_py.h
+++ b/paddle/fluid/pybind/tensor_py.h
@@ -190,6 +190,7 @@ void PyCUDATensorSetFromArray(
       static_cast<const platform::CUDADeviceContext *>(pool.Get(place));
   paddle::platform::GpuMemcpyAsync(dst, array.data(), sizeof(T) * array.size(),
                                    cudaMemcpyHostToDevice, dev_ctx->stream());
+  dev_ctx->Wait();
 }
 
 template <>