update: grad

Author: kozistr

pytorch_optimizer/optimizer/adashift.py

@@ -86,7 +86,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 state = self.state[p]
 
                 if len(state) == 0:
-                    state['grad_queue'] = deque([p.grad.clone()], maxlen=group['keep_num'])
+                    state['grad_queue'] = deque([grad.clone()], maxlen=group['keep_num'])
                     state['exp_avg'] = torch.zeros_like(p)
                     state['exp_avg_sq'] = torch.zeros_like(p)