optimize-offload support adamw op type (#36432)

PaddlePaddle · Oct 14, 2021 · 66c58fa · 66c58fa
1 parent 8256f6f
commit 66c58fa
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/paddle/distributed/fleet/meta_optimizers/sharding/offload_helper.py b/python/paddle/distributed/fleet/meta_optimizers/sharding/offload_helper.py
@@ -284,7 +284,7 @@ def offload(self, block, startup_block):
                 break
 
             vars_name = []
-            if op.type == "adam":
+            if op.type == "adam" or op.type == "adamw":
                 # {Moment1Out = [''], Moment2Out = [''], ParamOut = ['']} =
                 # adam(inputs={Moment1 = [''], Moment2 = [''], Param = ['']})
                 vars_name.append(op.desc.input("Moment1")[0])