Add ElementwiseOpInferVarType

PaddlePaddle · Mar 8, 2018 · 53d19f5 · 53d19f5
1 parent 0d49b92
commit 53d19f5
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/paddle/fluid/operators/elementwise_add_op.cc b/paddle/fluid/operators/elementwise_add_op.cc
@@ -29,8 +29,11 @@ class ElementwiseAddOpMaker : public ElementwiseOpMaker {
 }  // namespace paddle
 
 namespace ops = paddle::operators;
-REGISTER_OP(elementwise_add, ops::ElementwiseOp, ops::ElementwiseAddOpMaker,
-            elementwise_add_grad, ops::ElementwiseOpGrad);
+REGISTER_OPERATOR(elementwise_add, ops::ElementwiseOp,
+                  ops::ElementwiseAddOpMaker, ops::ElementwiseOpInferVarType,
+                  paddle::framework::DefaultGradOpDescMaker<true>);
+REGISTER_OPERATOR(elementwise_add_grad, ops::ElementwiseOpGrad);
+
 REGISTER_OP_CPU_KERNEL(
     elementwise_add,
     ops::ElementwiseAddKernel<paddle::platform::CPUDeviceContext, float>,

diff --git a/paddle/fluid/operators/elementwise_op.h b/paddle/fluid/operators/elementwise_op.h
@@ -41,6 +41,16 @@ class ElementwiseOp : public framework::OperatorWithKernel {
   }
 };
 
+class ElementwiseOpInferVarType : public framework::VarTypeInference {
+ public:
+  void operator()(const framework::OpDesc& op_desc,
+                  framework::BlockDesc* block) const override {
+    auto x_var = op_desc.Input("X")[0];
+    auto out_var = op_desc.Output("Out")[0];
+    block->Var(out_var)->SetType(block->Var(x_var)->GetType());
+  }
+};
+
 class ElementwiseOpMaker : public framework::OpProtoAndCheckerMaker {
  public:
   ElementwiseOpMaker(OpProto* proto, OpAttrChecker* op_checker)