fix: exp_avg_sq

Author: kozistr

pytorch_optimizer/optimizer/adopt.py

@@ -115,8 +115,6 @@ 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'])
 
                 normed_grad = grad.div(de_nom)
@@ -137,4 +135,6 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 p.add_(update, alpha=-lr)
 
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1.0 - beta2)
+
         return loss