1 files changed, 7 insertions, 0 deletions

diff --git a/train_lora.py b/train_lora.py
index d89b18d..59beb09 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -272,6 +272,12 @@ def parse_args():
         help="Number of steps for the warmup in the lr scheduler."
     )
     parser.add_argument(
+        "--lr_mid_point",
+        type=float,
+        default=0.3,
+        help="OneCycle schedule mid point."
+    )
+    parser.add_argument(
         "--lr_cycles",
         type=int,
         default=None,
@@ -662,6 +668,7 @@ def main():
         end_lr=1e2,
         train_epochs=num_train_epochs,
         warmup_epochs=args.lr_warmup_epochs,
+        mid_point=args.lr_mid_point,
     )
 
     metrics = trainer(