make one cycle lr div factor configurable

This commit is contained in:
Wing Lian
2023-05-21 20:25:06 -04:00
parent 0f74464652
commit 99383f14a3

View File

@@ -157,7 +157,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
cfg.learning_rate, cfg.learning_rate,
total_steps=total_num_steps, total_steps=total_num_steps,
epochs=cfg.num_epochs, epochs=cfg.num_epochs,
div_factor=10, div_factor=cfg.lr_div_factor if cfg.lr_div_factor else 6,
**lr_scheduler_kwargs, **lr_scheduler_kwargs,
) )
elif cfg.lr_scheduler == "log_sweep": elif cfg.lr_scheduler == "log_sweep":