Fixed LR finder

1 files changed, 2 insertions, 2 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 1e49474..218018b 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -634,7 +634,7 @@ def main():
         )
 
     if args.find_lr:
-        args.learning_rate = 1e2
+        args.learning_rate = 1e-4
 
     # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
     if args.use_8bit_adam:
@@ -901,7 +901,7 @@ def main():
             on_train=lambda: tokenizer.set_use_vector_shuffle(args.vector_shuffle),
             on_eval=lambda: tokenizer.set_use_vector_shuffle(False)
         )
-        lr_finder.run(min_lr=1e-4)
+        lr_finder.run(end_lr=1e2)
 
         plt.savefig(basepath.joinpath("lr.png"))
         plt.close()