refactor: nan_to_num()

Author: kozistr

pytorch_optimizer/optimizer/nero.py

@@ -92,7 +92,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 bias_correction: float = 1.0 - self.beta ** state['step']
 
                 grad_normed = grad / ((exp_avg_sq / bias_correction).sqrt() + self.eps)
-                grad_normed[torch.isnan(grad_normed)] = 0.0
+                torch.nan_to_num(grad_normed, nan=0.0, out=grad_normed)
 
                 p.sub_(group['lr'] * state['scale'] * grad_normed)