From 012e7e63a5b1acdf69c72eee4cb330a5a6defc41 Mon Sep 17 00:00:00 2001
From: Kohya S <ykumeykume@gmail.com>
Date: Sun, 29 Sep 2024 23:18:16 +0900
Subject: [PATCH] fix to work linear/cosine scheduler closes #1651 ref #1393

---
 library/train_util.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/library/train_util.py b/library/train_util.py
index 422dceca2..27910dc90 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -4496,6 +4496,15 @@ def wrap_check_needless_num_warmup_steps(return_vals):
             **lr_scheduler_kwargs,
         )
 
+    # these schedulers do not require `num_decay_steps`
+    if name == SchedulerType.LINEAR or name == SchedulerType.COSINE:
+        return schedule_func(
+            optimizer,
+            num_warmup_steps=num_warmup_steps,
+            num_training_steps=num_training_steps,
+            **lr_scheduler_kwargs,
+        )
+
     # All other schedulers require `num_decay_steps`
     if num_decay_steps is None:
         raise ValueError(f"{name} requires `num_decay_steps`, please provide that argument.")