fixed "mps" for load_hf_whisper()

-fixed `get_device()` returning `True` instead of "mps" when mps is available causing (#307)
jianfch · Feb 11, 2024 · 53272cb · 53272cb
1 parent 133f323
commit 53272cb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/stable_whisper/whisper_word_level/hf_whisper.py b/stable_whisper/whisper_word_level/hf_whisper.py
@@ -29,8 +29,8 @@ def get_device(device: str = None) -> str:
         return device
     if torch.cuda.is_available():
         return 'cuda:0'
-    if (mps := getattr(torch.backends, 'mps', None)) is not None:
-        return mps.is_available()
+    if (mps := getattr(torch.backends, 'mps', None)) is not None and mps.is_available():
+        return 'mps'
     return 'cpu'