diff --git a/paddle/operators/momentum_op.cc b/paddle/operators/momentum_op.cc
index efa0b59992fe6c4cfe057e6e631e17c3e7f494b7..9be4d15a43d87ae1a27c81498e8b19b0049a3bfa 100644
--- a/paddle/operators/momentum_op.cc
+++ b/paddle/operators/momentum_op.cc
@@ -22,7 +22,7 @@ class MomentumOp : public framework::OperatorWithKernel {
   using framework::OperatorWithKernel::OperatorWithKernel;
 
  protected:
-  void InferShape(framework::InferShapeContextBase *ctx) const override {
+  void InferShape(framework::InferShapeContext *ctx) const override {
     PADDLE_ENFORCE(ctx->HasInput("Param"),
                    "Input(param) of Momentum should not be null.");
     PADDLE_ENFORCE(ctx->HasInput("Grad"),