1 files changed, 14 insertions, 7 deletions

diff --git a/train_lora.py b/train_lora.py
index 54c9e7a..e81742a 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -81,6 +81,12 @@ def parse_args():
         help="The name of the current project.",
     )
     parser.add_argument(
+        "--auto_cycles",
+        type=int,
+        default=1,
+        help="How many cycles to run automatically."
+    )
+    parser.add_argument(
         "--placeholder_tokens",
         type=str,
         nargs='*',
@@ -933,10 +939,15 @@ def main():
         train_epochs=num_train_epochs,
     )
 
-    continue_training = True
-    training_iter = 1
+    training_iter = 0
+
+    while True:
+        training_iter += 1
+        if training_iter > args.auto_cycles:
+            response = input("Run another cycle? [y/n] ")
+            if response.lower().strip() == "n":
+                break
 
-    while continue_training:
         print("")
         print(f"============ LoRA cycle {training_iter} ============")
         print("")
@@ -961,10 +972,6 @@ def main():
             sample_frequency=lora_sample_frequency,
         )
 
-        response = input("Run another cycle? [y/n] ")
-        continue_training = response.lower().strip() != "n"
-        training_iter += 1
-
 
 if __name__ == "__main__":
     main()