diff --git a/paddle/fluid/operators/reshape_op.h b/paddle/fluid/operators/reshape_op.h index 807e5ad951b893a4c027a96d743f0606b70cf160..9abc78421a7554f51f56665e4d82d34e67c7c159 100644 --- a/paddle/fluid/operators/reshape_op.h +++ b/paddle/fluid/operators/reshape_op.h @@ -60,7 +60,7 @@ class ReshapeOp : public framework::OperatorWithKernel { static framework::DDim ValidateShape(const std::vector shape, const framework::DDim &in_dims) { const int64_t in_size = framework::product(in_dims); - // only one dimension canbe set to -1, whose size will be automatically + // only one dimension can be set to -1, whose size will be automatically // infered. const int64_t unk_dim_val = -1; const int64_t copy_dim_val = 0; @@ -119,13 +119,15 @@ class ReshapeKernel : public framework::OpKernel { auto *shape_tensor = ctx.Input("Shape"); framework::DDim out_dims = out->dims(); + if (shape_tensor) { auto *shape_data = shape_tensor->data(); + framework::Tensor cpu_shape_tensor; if (platform::is_gpu_place(ctx.GetPlace())) { - framework::Tensor cpu_shape_tensor; TensorCopy(*shape_tensor, platform::CPUPlace(), ctx.device_context(), &cpu_shape_tensor); shape_data = cpu_shape_tensor.data(); + ctx.device_context().Wait(); } auto shape = std::vector(shape_data, shape_data + shape_tensor->numel());