Fix

1 files changed, 3 insertions, 3 deletions

diff --git a/train_ti.py b/train_ti.py
index 880320f..b00b0d7 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -925,18 +925,18 @@ 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