update: codes

Author: kozistr

pytorch_optimizer/optimizer/experimental/ranger25.py

@@ -169,9 +169,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 state = self.state[p]
 
                 if len(state) == 0:
-                    state['exp_avg'] = torch.zeros_like(p)
-                    state['exp_avg_sq'] = torch.zeros_like(p)
-                    state['exp_avg_slow'] = torch.zeros_like(p)
+                    state['exp_avg'] = torch.zeros_like(grad)
+                    state['exp_avg_sq'] = torch.zeros_like(grad)
+                    state['exp_avg_slow'] = torch.zeros_like(grad)
                     state['slow_momentum'] = p.clone()
 
                 self.apply_weight_decay(