diff --git a/paddle/fluid/operators/gather_tree_op.cc b/paddle/fluid/operators/gather_tree_op.cc
index 26f9989121e9eb32d27817d23fa972e5e8fcf1b7..830134e57e0e72c5470ac79714015a94df9888bf 100644
--- a/paddle/fluid/operators/gather_tree_op.cc
+++ b/paddle/fluid/operators/gather_tree_op.cc
@@ -22,18 +22,16 @@ class GatherTreeOp : public framework::OperatorWithKernel {
   using framework::OperatorWithKernel::OperatorWithKernel;
 
   void InferShape(framework::InferShapeContext* ctx) const override {
-    PADDLE_ENFORCE(ctx->HasInput("Ids"),
-                   "Input(Ids) of GatherTreeOp should not be null.");
-    PADDLE_ENFORCE(ctx->HasInput("Parents"),
-                   "Input(Parents) of GatherTreeOp should not be null.");
-    PADDLE_ENFORCE(ctx->HasOutput("Out"),
-                   "Output(Out) of GatherTreeOp should not be null.");
+    OP_INOUT_CHECK(ctx->HasInput("Ids"), "Input", "Ids", "GatherTree");
+    OP_INOUT_CHECK(ctx->HasInput("Parents"), "Input", "Parents", "GatherTree");
+    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "GatherTree");
 
     auto ids_dims = ctx->GetInputDim("Ids");
     auto parents_dims = ctx->GetInputDim("Parents");
-    PADDLE_ENFORCE(ids_dims == parents_dims,
-                   "The shape of Input(Parents) must be same with the shape of "
-                   "Input(Ids).");
+    PADDLE_ENFORCE_EQ(ids_dims == parents_dims, true,
+                      platform::errors::InvalidArgument(
+                          "The shape of Input(Parents) must be same with the "
+                          "shape of Input(Ids)."));
     ctx->SetOutputDim("Out", ids_dims);
   }