fix: reset

kozistr · Aug 7, 2024 · e94290f · e94290f
1 parent 4f4d359
commit e94290f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pytorch_optimizer/optimizer/trac.py b/pytorch_optimizer/optimizer/trac.py
@@ -138,7 +138,7 @@ def state(self):
     def reset(self):
         device = self.param_groups[0]['params'][0].device
 
-        self.state = {
+        self.state['trac'] = {
             'betas': torch.tensor(self.betas, device=device),
             's': torch.zeros(len(self.betas), device=device),
             'variance': torch.zeros(len(self.betas), device=device),
@@ -148,7 +148,7 @@ def reset(self):
 
         for group in self.param_groups:
             for p in group['params']:
-                self.state[p] = p.clone()
+                self.state['trac'][p] = p.clone()
 
     @torch.no_grad()
     def zero_grad(self) -> None: