1 files changed, 2 insertions, 3 deletions

diff --git a/train_lora.py b/train_lora.py
index ba5aee1..d0313fe 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -967,18 +967,17 @@ def main():
         if len(auto_cycles) != 0:
             response = auto_cycles.pop(0)
         else:
-            response = input("Choose action: [o] one_cycle, [w] warmup, [c] constant, [d] decay, [s] stop \n--> ")
+            response = input("\n### Choose action: [o] one_cycle, [w] warmup, [c] constant, [d] decay, [s] stop \n--> ")
 
         if response.lower().strip() == "o":
             lr_scheduler = "one_cycle"
             lr_warmup_epochs = args.lr_warmup_epochs
             lr_cycles = args.lr_cycles
         if response.lower().strip() == "w":
-            lr_scheduler = "constant"
+            lr_scheduler = "constant_with_warmup"
             lr_warmup_epochs = num_train_epochs
         if response.lower().strip() == "c":
             lr_scheduler = "constant"
-            lr_warmup_epochs = 0
         if response.lower().strip() == "d":
             lr_scheduler = "cosine"
             lr_warmup_epochs = 0