未验证 提交 51c219c9 编写于 作者: F fengjiayi 提交者: GitHub

Merge pull request #9892 from JiayiFeng/refine_reshape_op

Add Wait() for reshape_op
...@@ -147,6 +147,7 @@ class ReshapeKernel : public framework::OpKernel<T> { ...@@ -147,6 +147,7 @@ class ReshapeKernel : public framework::OpKernel<T> {
if (!inplace) { if (!inplace) {
out->mutable_data<T>(ctx.GetPlace()); out->mutable_data<T>(ctx.GetPlace());
framework::TensorCopy(*in, ctx.GetPlace(), ctx.device_context(), out); framework::TensorCopy(*in, ctx.GetPlace(), ctx.device_context(), out);
ctx.device_context().Wait();
// TensorCopy will resize to in_dims. // TensorCopy will resize to in_dims.
out->Resize(out_dims); out->Resize(out_dims);
} else { } else {
...@@ -169,6 +170,7 @@ class ReshapeGradKernel : public framework::OpKernel<T> { ...@@ -169,6 +170,7 @@ class ReshapeGradKernel : public framework::OpKernel<T> {
auto in_dims = d_x->dims(); auto in_dims = d_x->dims();
if (!inplace) { if (!inplace) {
framework::TensorCopy(*d_out, ctx.GetPlace(), ctx.device_context(), d_x); framework::TensorCopy(*d_out, ctx.GetPlace(), ctx.device_context(), d_x);
ctx.device_context().Wait();
d_x->Resize(in_dims); d_x->Resize(in_dims);
} else { } else {
d_x->ShareDataWith(*d_out); d_x->ShareDataWith(*d_out);
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册