Fix condition scheduler
Browse files
src/axolotl/utils/trainer.py
CHANGED
@@ -104,7 +104,7 @@ def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer):
|
|
104 |
report_to="wandb" if cfg.use_wandb else None,
|
105 |
run_name=cfg.wandb_run_id if cfg.use_wandb else None,
|
106 |
optim=cfg.optimizer if cfg.optimizer else None,
|
107 |
-
lr_scheduler_type=cfg.lr_scheduler if cfg.lr_scheduler
|
108 |
weight_decay=cfg.weight_decay if cfg.weight_decay is not None else 0.0,
|
109 |
**training_arguments_kwargs,
|
110 |
)
|
|
|
104 |
report_to="wandb" if cfg.use_wandb else None,
|
105 |
run_name=cfg.wandb_run_id if cfg.use_wandb else None,
|
106 |
optim=cfg.optimizer if cfg.optimizer else None,
|
107 |
+
lr_scheduler_type=cfg.lr_scheduler if cfg.lr_scheduler in ("one_cycle", "log_sweep") else "cosine",
|
108 |
weight_decay=cfg.weight_decay if cfg.weight_decay is not None else 0.0,
|
109 |
**training_arguments_kwargs,
|
110 |
)
|