map only on one process (#13810)

2021-09-30 18:52:53 +02:00
parent 9a9805fccf
commit 44eb8bdeea
9 changed files with 142 additions and 125 deletions
--- a/examples/pytorch/multiple-choice/run_swag_no_trainer.py
+++ b/examples/pytorch/multiple-choice/run_swag_no_trainer.py
@@ -381,9 +381,10 @@ def main():
        tokenized_inputs["labels"] = labels
        return tokenized_inputs

-    processed_datasets = raw_datasets.map(
-        preprocess_function, batched=True, remove_columns=raw_datasets["train"].column_names
-    )
+    with accelerator.main_process_first():
+        processed_datasets = raw_datasets.map(
+            preprocess_function, batched=True, remove_columns=raw_datasets["train"].column_names
+        )

    train_dataset = processed_datasets["train"]
    eval_dataset = processed_datasets["validation"]