Pass device to pyannote.audio.Inference

m-bain · Feb 21, 2023 · 57f5957 · 57f5957
1 parent f7093e6
commit 57f5957
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/whisperx/transcribe.py b/whisperx/transcribe.py
@@ -645,9 +645,12 @@ def cli():
         if hf_token is None:
             print("Warning, no huggingface token used, needs to be saved in environment variable, otherwise will throw error loading VAD model...")
         from pyannote.audio import Inference
-        vad_pipeline = Inference("pyannote/segmentation",
-                            pre_aggregation_hook=lambda segmentation: segmentation,
-                            use_auth_token=hf_token)
+        vad_pipeline = Inference(
+            "pyannote/segmentation",
+            pre_aggregation_hook=lambda segmentation: segmentation,
+            use_auth_token=hf_token,
+            device=torch.device(device),
+        )
 
     diarize_pipeline = None
     if diarize: