1 files changed, 3 insertions, 3 deletions

diff --git a/dreambooth_plus.py b/dreambooth_plus.py
index fa3a22b..06ff45b 100644
--- a/dreambooth_plus.py
+++ b/dreambooth_plus.py
@@ -125,7 +125,7 @@ def parse_args():
     parser.add_argument(
         "--max_train_steps",
         type=int,
-        default=1400,
+        default=2400,
         help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
     )
     parser.add_argument(
@@ -752,8 +752,8 @@ def main():
             optimizer=optimizer,
             num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
             num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
-            num_cycles=args.lr_cycles or math.ceil(
-                ((args.max_train_steps - args.lr_warmup_steps) / num_update_steps_per_epoch) / 2),
+            num_cycles=args.lr_cycles or math.ceil(math.sqrt(
+                ((args.max_train_steps - args.lr_warmup_steps) / num_update_steps_per_epoch))),
         )
     else:
         lr_scheduler = get_scheduler(