fix train script for babyai

huggingface · Oct 16, 2023 · d1c4bc9 · d1c4bc9
1 parent 9288371
commit d1c4bc9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scripts/train_gia2.py b/scripts/train_gia2.py
@@ -145,7 +145,7 @@ def main():
                 lambda example_batch: processor(**example_batch, padding="max_length", truncation="preserve"),
                 batched=True,
                 batch_size=10,
-                remove_columns={"text", "images"}.intersection(column_names),
+                remove_columns={"text", "images", "text_observations"}.intersection(column_names),
             )
             dataset = dataset.map(
                 lambda x: {"loss_weight": [LOSS_WEIGHTS.get(task, 1.0)] * len(next(iter(x.values())))}