From 922dee3b436194ee5cdde4eec1413f9c55281ed0 Mon Sep 17 00:00:00 2001 From: Yu Yang Date: Tue, 18 Sep 2018 07:21:05 +0000 Subject: [PATCH] Wait input when data transform --- paddle/fluid/framework/data_device_transform.cc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/paddle/fluid/framework/data_device_transform.cc b/paddle/fluid/framework/data_device_transform.cc index 6bcfc6cd55..fee6ba4004 100644 --- a/paddle/fluid/framework/data_device_transform.cc +++ b/paddle/fluid/framework/data_device_transform.cc @@ -25,6 +25,10 @@ void TransDataDevice(const Tensor &in, const platform::Place &dst_place, in.place().which(), dst_place.which(), "Currently, model parallelism is only supported between CPU and CUDA"); + // NOTE(yy): TransDataDevice should wait for computation of input. + platform::DeviceContextPool::Instance().Get(in.place())->Wait(); + platform::DeviceContextPool::Instance().Get(dst_place)->Wait(); + // FIXME(zcd): TransDataDevice is used to transform data from GPU to CPU and // the enforced checkings have been done in GetDeviceContext, so the // `dev_ctx->Wait()` is necessary. But `dev_ctx->Wait()` will make the program -- GitLab