diff --git a/paddle/fluid/operators/size_op.cc b/paddle/fluid/operators/size_op.cc index b45fa7c791ff22be422ce12a8348a071c60ddd0f..70733d643673ad8acde9a45f273a52a9723fb0d3 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 3ea3032693236d5618ff6f0c858cbd85e34633ab..de56ecd95270577689f699462b9273b43f34595e 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 fb44070897156ef88062231322e28a2db1f244a7..e8c53d6e683305bfc1ff7c052a2dc54ecf465936 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 363c3ffceb85ef6168dc8c33b81185cac08083fb..71ac809ddf70e071a9d409093a699b5d666af934 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]