Fixed LR finder

1 files changed, 3 insertions, 3 deletions

diff --git a/train_ti.py b/train_ti.py
index 2b3f017..102c0fa 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -584,7 +584,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:
@@ -853,9 +853,9 @@ 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.savefig(basepath.joinpath("lr.png"), dpi=300)
         plt.close()
 
         quit()