diff --git a/paddle/fluid/operators/optimizers/momentum_op.cc b/paddle/fluid/operators/optimizers/momentum_op.cc
index 1b01f5ebd879f265360db5723a3adc7cc3ece63d..b9a74c1bf7124e6702745685eff08b12d2ecb85e 100644
--- a/paddle/fluid/operators/optimizers/momentum_op.cc
+++ b/paddle/fluid/operators/optimizers/momentum_op.cc
@@ -118,6 +118,10 @@ REGISTER_OP_VERSION(momentum)
       Upgrade momentum add 2 attributes [regularization_method, regularization_coeff].
     )ROC",
         paddle::framework::compatible::OpVersionDesc()
+            .NewInput("MasterParam", "FP32 master weight for AMP.")
+            .NewOutput("MasterParamOut",
+                       "The updated FP32 master weight for AMP. "
+                       "It shared memory with Input(MasterParam).")
             .NewAttr("regularization_method",
                      "(string) regularization_method, right now only support "
                      "l2decay or none",