From 411df76a97dfe34c84c182ce7908eeeb38553404 Mon Sep 17 00:00:00 2001 From: Dan Saunders Date: Mon, 17 Mar 2025 22:57:55 +0000 Subject: [PATCH] bugfix --- src/axolotl/core/trainers/base.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/axolotl/core/trainers/base.py b/src/axolotl/core/trainers/base.py index 70f3a86ea..d81ad6d41 100644 --- a/src/axolotl/core/trainers/base.py +++ b/src/axolotl/core/trainers/base.py @@ -613,7 +613,8 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, Trainer): self.data_collator = ( # pylint: disable=attribute-defined-outside-init self.eval_data_collator ) - eval_dataset = eval_dataset.remove_columns(["length"]) + if "length" in eval_dataset.column_names: + eval_dataset = eval_dataset.remove_columns(["length"]) dataloader = super().get_eval_dataloader(eval_dataset) self.data_collator = ( # pylint: disable=attribute-defined-outside-init self.train_data_collator @@ -633,7 +634,8 @@ class AxolotlTrainer(SchedulerMixin, OptimizerMixin, Trainer): if hasattr(self, "eval_data_collator") and self.eval_data_collator else self.data_collator ) - eval_dataset = eval_dataset.remove_columns(["length"]) + if "length" in eval_dataset.column_names: + eval_dataset = eval_dataset.remove_columns(["length"]) # Handle dataset preprocessing for SP if self.args.sequence_parallel_degree > 1: