clean up InvSqrtWithWarmup

allenai · Aug 3, 2023 · c157fca · c157fca
1 parent e167b7d
commit c157fca
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/olmo/optim.py b/olmo/optim.py
@@ -109,7 +109,7 @@ def get_lr(self, initial_lr: float, step: int, max_steps: int) -> float:
         if step < self.warmup_steps:
             return self._linear_warmup(initial_lr, step, self.warmup_steps)
         del max_steps
-        return initial_lr * sqrt(self.warmup_steps) / sqrt(max(self.warmup_steps, step))
+        return initial_lr * sqrt(self.warmup_steps / max(self.warmup_steps, step))
 
 
 @dataclass