From c20cc2bd8a018f078e3916e01579df8faab66f92 Mon Sep 17 00:00:00 2001
From: fengjiayi <fengjiayi@baidu.com>
Date: Fri, 13 Apr 2018 05:45:52 +0000
Subject: [PATCH] Add Wait() for reshape_op

---
 paddle/fluid/operators/reshape_op.h | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/paddle/fluid/operators/reshape_op.h b/paddle/fluid/operators/reshape_op.h
index 9abc78421a7..8320c257c9a 100644
--- a/paddle/fluid/operators/reshape_op.h
+++ b/paddle/fluid/operators/reshape_op.h
@@ -147,6 +147,7 @@ class ReshapeKernel : public framework::OpKernel<T> {
     if (!inplace) {
       out->mutable_data<T>(ctx.GetPlace());
       framework::TensorCopy(*in, ctx.GetPlace(), ctx.device_context(), out);
+      ctx.device_context().Wait();
       // TensorCopy will resize to in_dims.
       out->Resize(out_dims);
     } else {
@@ -169,6 +170,7 @@ class ReshapeGradKernel : public framework::OpKernel<T> {
     auto in_dims = d_x->dims();
     if (!inplace) {
       framework::TensorCopy(*d_out, ctx.GetPlace(), ctx.device_context(), d_x);
+      ctx.device_context().Wait();
       d_x->Resize(in_dims);
     } else {
       d_x->ShareDataWith(*d_out);
-- 
GitLab