1 files changed, 9 insertions, 0 deletions

diff --git a/train_ti.py b/train_ti.py
index 45e730a..c452269 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -79,6 +79,11 @@ def parse_args():
         help="Learning rate decay per cycle."
     )
     parser.add_argument(
+        "--cycle_constant",
+        action="store_true",
+        help="Use constant LR on cycles > 1."
+    )
+    parser.add_argument(
         "--placeholder_tokens",
         type=str,
         nargs='*',
@@ -926,6 +931,10 @@ def main():
             print(f"------------ TI cycle {training_iter + 1} ------------")
             print("")
 
+            if args.cycle_constant and training_iter == 1:
+                args.lr_scheduler = "constant"
+                args.lr_warmup_epochs = 0
+
             optimizer = create_optimizer(
                 text_encoder.text_model.embeddings.token_embedding.parameters(),
                 lr=learning_rate,