update: optimize the code

kozistr · kozistr · commit 13d9565c1a96 · 2025-12-21T03:57:28.000+09:00
diff --git a/pytorch_optimizer/optimizer/ano.py b/pytorch_optimizer/optimizer/ano.py
@@ -121,14 +121,12 @@ def step(self, closure: Closure = None) -> Loss:
                 exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
 
                 square_grad = grad.square()
-                sign_term = torch.sign(square_grad - exp_avg_sq)
-                exp_avg_sq.mul_(beta2).add_(sign_term * square_grad, alpha=1.0 - beta2)
-
-                v_hat = exp_avg_sq / bias_correction2
-                adjusted_lr = group['lr'] / v_hat.sqrt_().add_(group['eps'])
+                exp_avg_sq.mul_(beta2).addcmul_(
+                    torch.sign(square_grad - exp_avg_sq), square_grad, value=1.0 - beta2
+                )
 
-                update = adjusted_lr * grad.abs() * exp_avg.sign()
+                de_nom = square_grad.copy_(exp_avg_sq).div_(bias_correction2).sqrt_().add_(group['eps'])
 
-                p.add_(update, alpha=-1.0)
+                p.addcdiv_(grad.abs().mul_(exp_avg.sign()), de_nom, value=-group['lr'])
 
         return loss
