Update

1 files changed, 14 insertions, 3 deletions

diff --git a/train_ti.py b/train_ti.py
index ca5b113..ebac302 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -64,6 +64,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='*',
@@ -869,10 +875,15 @@ def main():
             mid_point=args.lr_mid_point,
         )
 
-        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"------------ TI cycle {training_iter} ------------")
             print("")