1 files changed, 9 insertions, 4 deletions

diff --git a/train_ti.py b/train_ti.py
index b00b0d7..84ca296 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -906,9 +906,6 @@ def main():
         if args.sample_num is not None:
             sample_frequency = math.ceil(num_train_epochs / args.sample_num)
 
-        training_iter = 0
-        learning_rate = args.learning_rate
-
         project = placeholder_tokens[0] if len(placeholder_tokens) == 1 else "ti"
 
         if accelerator.is_main_process:
@@ -916,7 +913,9 @@ def main():
 
         sample_output_dir = output_dir / project / "samples"
 
+        training_iter = 0
         auto_cycles = list(args.auto_cycles)
+        learning_rate = args.learning_rate
         lr_scheduler = args.lr_scheduler
         lr_warmup_epochs = args.lr_warmup_epochs
         lr_cycles = args.lr_cycles
@@ -929,6 +928,12 @@ def main():
                     "\n### Choose action: [o] one_cycle, [w] warmup, [c] constant, [d] decay, [s] stop \n--> ")
 
             if response.lower().strip() == "o":
+                if args.learning_rate is not None:
+                    learning_rate = args.learning_rate * 2
+            else:
+                learning_rate = args.learning_rate
+
+            if response.lower().strip() == "o":
                 lr_scheduler = "one_cycle"
                 lr_warmup_epochs = args.lr_warmup_epochs
                 lr_cycles = args.lr_cycles
@@ -945,7 +950,7 @@ def main():
                 break
 
             print("")
-            print(f"------------ TI cycle {training_iter + 1} ------------")
+            print(f"------------ TI cycle {training_iter + 1}: {response} ------------")
             print("")
 
             optimizer = create_optimizer(