From d7d49f93158459f4abbae1242ae56c77d1203627 Mon Sep 17 00:00:00 2001
From: Mansheej Paul <mansheej.paul@databricks.com>
Date: Tue, 10 Oct 2023 21:59:25 +0000
Subject: [PATCH] more fixes

---
 llmfoundry/optim/scheduler.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/llmfoundry/optim/scheduler.py b/llmfoundry/optim/scheduler.py
index 6e863a5467..5043c1d2d8 100644
--- a/llmfoundry/optim/scheduler.py
+++ b/llmfoundry/optim/scheduler.py
@@ -87,9 +87,9 @@ def __init__(self,
                  alpha_f_cooldown: float = 0.0,
                  scale_warmup: bool = False) -> None:
         if alpha_f_decay < alpha_f_cooldown:
-            raise ValueError(
-                f'Required: alpha_f_decay >= alpha_f_cooldown. Current: alpha_f_decay={alpha_f_decay}, alpha_f_cooldown={alpha_f_cooldown}'
-            )
+            raise ValueError(('Required: alpha_f_decay >= alpha_f_cooldown. '
+                              f'Current: alpha_f_decay={alpha_f_decay}, '
+                              f'alpha_f_cooldown={alpha_f_cooldown}.'))
         _raise_if_units_dur(t_warmup, 't_warmup')
         _raise_if_units_dur(t_scale, 't_scale')
         _raise_if_units_dur(t_cooldown, 't_cooldown')
@@ -131,6 +131,7 @@ def __call__(self, state: State, ssr: float = 1.0) -> float:
         current_time = state.timestamp.get(t_scale.unit)
 
         t_shift = t_scale - t_warmup
+        # t_cooldown_start = max(t_warmup, t_max - t_cooldown)
         t_cooldown_start = t_max - t_cooldown
         if t_cooldown_start < t_warmup:
             t_cooldown_start = t_warmup