1 files changed, 7 insertions, 0 deletions

diff --git a/train_ti.py b/train_ti.py
index b182a72..83043ad 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -258,6 +258,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,
@@ -790,6 +796,7 @@ def main():
             end_lr=1e3,
             train_epochs=num_train_epochs,
             warmup_epochs=args.lr_warmup_epochs,
+            mid_point=args.lr_mid_point,
         )
 
         metrics = trainer(