Fix remove_columns in text-classification example (#29351)

huggingface · Apr 22, 2024 · ad714a9 · ad714a9
1 parent 23b1cbf
commit ad714a9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/pytorch/text-classification/run_classification.py b/examples/pytorch/text-classification/run_classification.py
@@ -422,7 +422,7 @@ def main():
         for split in raw_datasets.keys():
             for column in data_args.remove_columns.split(","):
                 logger.info(f"removing column {column} from split {split}")
-                raw_datasets[split].remove_columns(column)
+                raw_datasets[split] = raw_datasets[split].remove_columns(column)
 
     if data_args.label_column_name is not None and data_args.label_column_name != "label":
         for key in raw_datasets.keys():