Update pytorch_optimizer/optimizer/soap.py

Author: kozistr

pytorch_optimizer/optimizer/soap.py

@@ -269,8 +269,8 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 state = self.state[p]
                 if len(state) == 0:
-                    state['exp_avg'] = torch.zeros_like(grad, dtype=p.dtype)
-                    state['exp_avg_sq'] = torch.zeros_like(grad, dtype=p.dtype)
+                    state['exp_avg'] = torch.zeros_like(grad)
+                    state['exp_avg_sq'] = torch.zeros_like(grad)
 
                     self.init_pre_conditioner(
                         grad,