Adding control point for scale adjustment

111ee132 · Michael Carilli · d69011de · 111ee132
Commit 111ee132 authored Apr 16, 2019 by Michael Carilli
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

apex/optimizers/fused_adam.py apex/optimizers/fused_adam.py +1 -1

No files found.
--- a/apex/optimizers/fused_adam.py
+++ b/apex/optimizers/fused_adam.py
@@ -81,7 +81,7 @@ class FusedAdam(torch.optim.Optimizer):
        if hasattr(self, "_amp_stash"):
            grads = self._amp_stash.grads
            output_params = self._amp_stash.output_params
-            scale = self._amp_stash.scale
+            scale = self._amp_stash.scale*scale
            grad_norms = self._amp_stash.grad_norms

        if grads is None: