fix momentum op register (#29941)

* fix momentum op register
4 years ago · 4cbcc9b6da
parent 7c1f69bdf0
commit 4cbcc9b6da
1 changed files with 4 additions and 0 deletions
--- 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",