fix: lerp

kozistr · May 5, 2024 · 0f2d6b7 · 0f2d6b7
1 parent 518e32d
commit 0f2d6b7
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/pytorch_optimizer/optimizer/schedulefree.py b/pytorch_optimizer/optimizer/schedulefree.py
@@ -71,7 +71,7 @@ def eval(self):
                 for p in group['params']:
                     state = self.state[p]
                     if 'z' in state:
-                        p.lerp_(end=state['z'], weight=1.0 - 1.0 / momentum)
+                        p.data.lerp_(end=state['z'], weight=1.0 - 1.0 / momentum)
                 group['train_mode'] = False
 
     def train(self):
@@ -81,7 +81,7 @@ def train(self):
                 for p in group['params']:
                     state = self.state[p]
                     if 'z' in state:
-                        p.lerp_(end=state['z'], weight=1.0 - momentum)
+                        p.data.lerp_(end=state['z'], weight=1.0 - momentum)
                 group['train_mode'] = True
 
     @torch.no_grad()
@@ -216,7 +216,7 @@ def eval(self):
                 for p in group['params']:
                     state = self.state[p]
                     if 'z' in state:
-                        p.lerp_(end=state['z'], weight=1.0 - 1.0 / beta1)
+                        p.data.lerp_(end=state['z'], weight=1.0 - 1.0 / beta1)
                 group['train_mode'] = False
 
     def train(self):
@@ -226,7 +226,7 @@ def train(self):
                 for p in group['params']:
                     state = self.state[p]
                     if 'z' in state:
-                        p.lerp_(end=state['z'], weight=1.0 - beta1)
+                        p.data.lerp_(end=state['z'], weight=1.0 - beta1)
                 group['train_mode'] = True
 
     @torch.no_grad()