make one cycle lr div factor configurable

2023-05-21 20:25:06 -04:00
parent 0f74464652
commit 99383f14a3
1 changed files with 2 additions and 2 deletions
--- a/src/axolotl/utils/trainer.py
+++ b/src/axolotl/utils/trainer.py
@@ -157,7 +157,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
                cfg.learning_rate,
                total_steps=total_num_steps,
                epochs=cfg.num_epochs,
-                div_factor=10,
+                div_factor=cfg.lr_div_factor if cfg.lr_div_factor else 6,
                **lr_scheduler_kwargs,
            )
        elif cfg.lr_scheduler == "log_sweep":