From 066949533d559b42fb06fe37b1b78dcb72bc823a Mon Sep 17 00:00:00 2001 From: sneaxiy <32832641+sneaxiy@users.noreply.github.com> Date: Tue, 26 Apr 2022 15:59:34 +0800 Subject: [PATCH] Make size op do no data transformation (#42204) * fix size op * update --- paddle/fluid/operators/size_op.cc | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/paddle/fluid/operators/size_op.cc b/paddle/fluid/operators/size_op.cc index 84b0f403be0..4af355bfca6 100644 --- a/paddle/fluid/operators/size_op.cc +++ b/paddle/fluid/operators/size_op.cc @@ -23,6 +23,19 @@ namespace operators { class SizeOp : public framework::OperatorWithKernel { public: using framework::OperatorWithKernel::OperatorWithKernel; + + protected: + framework::OpKernelType GetExpectedKernelType( + const framework::ExecutionContext& ctx) const override { + auto dtype = framework::proto::VarType::FP32; // dtype is not important + return framework::OpKernelType(dtype, ctx.GetPlace()); + } + + framework::OpKernelType GetKernelTypeForVar( + const std::string& var_name, const framework::Tensor& tensor, + const framework::OpKernelType& expected_kernel_type) const override { + return expected_kernel_type; + } }; class SizeOpMaker : public framework::OpProtoAndCheckerMaker { @@ -40,6 +53,8 @@ Return the number of elements in the input. } }; +DECLARE_NO_NEED_BUFFER_VARS_INFERER(SizeOpNoNeedBufferVarInferer, "Input"); + } // namespace operators } // namespace paddle @@ -50,4 +65,4 @@ REGISTER_OPERATOR( size, ops::SizeOp, ops::SizeOpMaker, paddle::framework::EmptyGradOpMaker, paddle::framework::EmptyGradOpMaker, - SizeInferShapeFunctor); + SizeInferShapeFunctor, ops::SizeOpNoNeedBufferVarInferer); -- GitLab