From fce52cefa744a5900fc065aafb2f55d846d1202c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20=C5=A0a=C5=A1ko?= Date: Tue, 2 Apr 2024 19:15:27 +0200 Subject: [PATCH] Fix `remove_columns` in `text-classification` example (#29351) --- examples/pytorch/text-classification/run_classification.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/pytorch/text-classification/run_classification.py b/examples/pytorch/text-classification/run_classification.py index 0b3d6517c7..982dbf9cc7 100755 --- 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():