fix(eval): implement DataParallel for multi-GPU inferencing

huggingface · Sep 12, 2024 · 5b1f181 · 5b1f181
1 parent 5818ca3
commit 5b1f181
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/training/run_eval.py b/training/run_eval.py
@@ -47,6 +47,7 @@
 from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE
 from transformers.utils import check_min_version, is_accelerate_available
 from transformers.utils.versions import require_version
+from torch.nn import DataParallel
 
 
 # Will error if the minimal version of Transformers is not installed. Remove at your own risks.
@@ -561,6 +562,9 @@ def main():
         cache_dir=data_args.cache_dir,
         variant=data_args.model_variant,
     )
+    if torch.cuda.device_count() > 1:
+        print(f"Using {torch.cuda.device_count()} GPUs")
+        model = DataParallel(model)
     model.to("cuda:0", dtype=dtype)
 
     model_pipeline = None