From 4eeda9d67616ec57d01bb436d4d60dd1a1fde458 Mon Sep 17 00:00:00 2001
From: Zeng Jinle <32832641+sneaxiy@users.noreply.github.com>
Date: Fri, 18 Oct 2019 10:05:54 +0800
Subject: [PATCH] fix tensor_util, test=develop (#20699)

---
 paddle/fluid/framework/tensor_util_test.cc | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/paddle/fluid/framework/tensor_util_test.cc b/paddle/fluid/framework/tensor_util_test.cc
index bf9dabcd23..ef9964b3ae 100644
--- a/paddle/fluid/framework/tensor_util_test.cc
+++ b/paddle/fluid/framework/tensor_util_test.cc
@@ -281,12 +281,15 @@ TEST(TensorFromDLPack, Tensor) {
     // Copy to GPUTensor
     gpu_tensor.Resize(paddle::framework::make_ddim({3, 3}));
     paddle::platform::CUDAPlace gpu_place;
-    paddle::platform::CUDADeviceContext gpu_ctx(gpu_place);
+    auto& gpu_ctx =
+        *paddle::platform::DeviceContextPool::Instance().GetByPlace(gpu_place);
     paddle::framework::TensorFromVector<int>(src_vec, gpu_ctx, &gpu_tensor);
+    gpu_ctx.Wait();
 
     paddle::framework::DLPackTensor dlpack_tensor(gpu_tensor, 1);
-
     paddle::framework::TensorFromDLPack(dlpack_tensor, &gpu_tensor_from_dlpack);
+    gpu_ctx.Wait();
+
     // Copy from GPU to CPU tensor for comparison
     paddle::framework::TensorCopy(gpu_tensor_from_dlpack, cpu_place, gpu_ctx,
                                   &dst_tensor);
-- 
GitLab