From ad229ffa919841e4b72da7479f4c56b519f8ed10 Mon Sep 17 00:00:00 2001 From: NanoCode012 Date: Fri, 16 May 2025 13:36:11 +0700 Subject: [PATCH] fix: add missing rex from rebase --- src/axolotl/core/trainer_builder/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/axolotl/core/trainer_builder/base.py b/src/axolotl/core/trainer_builder/base.py index 4364449c5..878718199 100644 --- a/src/axolotl/core/trainer_builder/base.py +++ b/src/axolotl/core/trainer_builder/base.py @@ -326,7 +326,7 @@ class TrainerBuilderBase(abc.ABC): # optim/scheduler training_args_kwargs["loraplus_lr_ratio"] = self.cfg.loraplus_lr_ratio training_args_kwargs["loraplus_lr_embedding"] = self.cfg.loraplus_lr_embedding - if self.cfg.lr_scheduler in ["one_cycle", "log_sweep"]: + if self.cfg.lr_scheduler in ["one_cycle", "log_sweep", "rex"]: training_args_kwargs["lr_scheduler_type"] = "cosine" training_args_kwargs["alternate_lr_scheduler_type"] = self.cfg.lr_scheduler else: