From 6ec7368f0c8a7102c5308b2d8d0672aecadbf978 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 23 Jan 2024 09:32:46 -0500 Subject: [PATCH] check for length before trying to remove it --- src/axolotl/core/trainer_builder.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index e109db7f84..1a145e3584 100644 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -227,7 +227,8 @@ def _get_eval_sampler( def get_train_dataloader(self) -> DataLoader: if self.args.sample_packing and not self.args.pretraining: train_dataset = self.train_dataset - train_dataset = train_dataset.remove_columns(["length"]) + if "length" in train_dataset.features.keys(): + train_dataset = train_dataset.remove_columns(["length"]) data_collator = self.data_collator dataloader_params = { "batch_size": self._train_batch_size,