refactor: order

kozistr · Nov 9, 2024 · 26d77ed · 26d77ed
1 parent e65c0cd
commit 26d77ed
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pytorch_optimizer/optimizer/adopt.py b/pytorch_optimizer/optimizer/adopt.py
@@ -104,13 +104,14 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     exp_avg_sq.addcmul_(grad, grad.conj())
                     continue
 
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1.0 - beta2)
+
                 de_nom = exp_avg_sq.sqrt().clamp_(min=group['eps'])
                 if group['step'] == 2:
                     exp_avg.addcdiv_(grad, de_nom)
                 else:
                     exp_avg.mul_(beta1).addcdiv_(grad, de_nom, value=1.0 - beta1)
 
                 p.add_(exp_avg, alpha=-group['lr'])
-                exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1.0 - beta2)
 
         return loss