From 613f238e56fe2e31af9657fc8eadeff08444591f Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 6 Nov 2024 09:43:35 -0500 Subject: [PATCH] use kwargs to support patch release --- src/axolotl/core/trainer_builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index f4a2f9001..4fadd7eb4 100755 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -896,13 +896,13 @@ class AxolotlTrainer(SchedulerMixin, Trainer): for key, value in metrics.items(): self._stored_metrics[train_eval][key].append(value) - def _save_checkpoint(self, model, trial): + def _save_checkpoint(self, model, trial, **kwargs): # make sure the checkpoint dir exists, since trainer is flakey checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}" run_dir = self._get_output_dir(trial=trial) output_dir = os.path.join(run_dir, checkpoint_folder) os.makedirs(output_dir, exist_ok=True) - return super()._save_checkpoint(model, trial) + return super()._save_checkpoint(model, trial, **kwargs) class AxolotlMambaTrainer(AxolotlTrainer):