Merge pull request #13 from ColoredCow/feature/make-sure-using-cuda

disable whishper modes
ColoredCow · Nov 12, 2024 · b82a8af · b82a8af
2 parents e3ba829 + 5f651e8
commit b82a8af
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/transcription.py b/transcription.py
@@ -8,6 +8,8 @@
 def load_asr_model(modelName):
     processor = WhisperProcessor.from_pretrained(modelName)
     model = WhisperForConditionalGeneration.from_pretrained(modelName)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device)
     return processor, model
 
 def transcribe_audio(file_path, model, processor, language):