Fix cycle loop

1 files changed, 5 insertions, 3 deletions

diff --git a/train_ti.py b/train_ti.py
index 84ca296..d1e5467 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -937,17 +937,19 @@ def main():
                 lr_scheduler = "one_cycle"
                 lr_warmup_epochs = args.lr_warmup_epochs
                 lr_cycles = args.lr_cycles
-            if response.lower().strip() == "w":
+            elif response.lower().strip() == "w":
                 lr_scheduler = "constant_with_warmup"
                 lr_warmup_epochs = num_train_epochs
-            if response.lower().strip() == "c":
+            elif response.lower().strip() == "c":
                 lr_scheduler = "constant"
-            if response.lower().strip() == "d":
+            elif response.lower().strip() == "d":
                 lr_scheduler = "cosine"
                 lr_warmup_epochs = 0
                 lr_cycles = 1
             elif response.lower().strip() == "s":
                 break
+            else:
+                continue
 
             print("")
             print(f"------------ TI cycle {training_iter + 1}: {response} ------------")