1 files changed, 7 insertions, 0 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 0634376..2c884d2 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -240,6 +240,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,
@@ -634,6 +640,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(