check for length before trying to remove it

axolotl-ai-cloud · Jan 23, 2024 · 6ec7368 · 6ec7368
1 parent fb7f9b9
commit 6ec7368
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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,