diff --git a/paddle/fluid/operators/size_op.cc b/paddle/fluid/operators/size_op.cc index b45fa7c791ff2..70733d643673a 100644 --- a/paddle/fluid/operators/size_op.cc +++ b/paddle/fluid/operators/size_op.cc @@ -53,7 +53,7 @@ REGISTER_OPERATOR( size, ops::SizeOp, ops::SizeOpMaker, paddle::framework::EmptyGradOpMaker, paddle::framework::EmptyGradOpMaker); -REGISTER_OP_CPU_KERNEL(size, ops::SizeKernel, ops::SizeKernel, +REGISTER_OP_CPU_KERNEL(size, ops::SizeKernel, ops::SizeKernel, ops::SizeKernel, ops::SizeKernel, ops::SizeKernel, ops::SizeKernel); diff --git a/paddle/fluid/operators/size_op.cu b/paddle/fluid/operators/size_op.cu index 3ea3032693236..de56ecd952705 100644 --- a/paddle/fluid/operators/size_op.cu +++ b/paddle/fluid/operators/size_op.cu @@ -16,7 +16,7 @@ limitations under the License. */ REGISTER_OP_CUDA_KERNEL( size, paddle::operators::SizeKernel, - paddle::operators::SizeKernel, + paddle::operators::SizeKernel, paddle::operators::SizeKernel, paddle::operators::SizeKernel, paddle::operators::SizeKernel, paddle::operators::SizeKernel); diff --git a/paddle/fluid/operators/size_op.h b/paddle/fluid/operators/size_op.h index fb44070897156..e8c53d6e68330 100644 --- a/paddle/fluid/operators/size_op.h +++ b/paddle/fluid/operators/size_op.h @@ -26,8 +26,18 @@ class SizeKernel : public framework::OpKernel { void Compute(const framework::ExecutionContext& ctx) const override { auto* in_t = ctx.Input("Input"); auto* out_t = ctx.Output("Out"); - auto out_data = out_t->mutable_data(platform::CPUPlace()); - out_data[0] = in_t->numel(); + auto place = ctx.GetPlace(); + auto out_data = out_t->mutable_data(place); + auto cpu_place = platform::CPUPlace(); + if (place == cpu_place) { + out_data[0] = in_t->numel(); + } else { + Tensor cpu_tensor; + auto cpu_data = + cpu_tensor.mutable_data(out_t->dims(), cpu_place); + cpu_data[0] = in_t->numel(); + TensorCopy(cpu_tensor, place, out_t); + } } }; } // namespace operators diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py index 363c3ffceb85e..71ac809ddf70e 100644 --- a/python/paddle/tensor/manipulation.py +++ b/python/paddle/tensor/manipulation.py @@ -1001,7 +1001,7 @@ def chunk(x, chunks, axis=0, name=None): x_np = np.random.random([3, 9, 5]).astype("int32") x = paddle.to_tensor(x_np) - out0, out1, out22 = paddle.chunk(x, chunks=3, axis=1) + out0, out1, out2 = paddle.chunk(x, chunks=3, axis=1) # out0.shape [3, 3, 5] # out1.shape [3, 3, 5] # out2.shape [3, 3, 5]