Update src/tools/accent_tool.py

src/tools/accent_tool.py

@@ -5,6 +5,9 @@ import subprocess
 from pydub import AudioSegment
 import whisper
 from speechbrain.pretrained.interfaces import foreign_class
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def clear_tmp_dir(path):
     for filename in os.listdir(path):
@@ -19,12 +22,13 @@ def clear_tmp_dir(path):
 
 class AccentAnalyzerTool:
     def __init__(self):
-        self.whisper_model = whisper.load_model("medium")
+        self.whisper_model = whisper.load_model("medium", device = device)
         self.accent_model = foreign_class(
             source="Jzuluaga/accent-id-commonaccent_xlsr-en-english",
             pymodule_file="custom_interface.py",
             classname="CustomEncoderWav2vec2Classifier"
         )
+        self.accent_model.device = torch.device(device)
         self.last_transcript = None
 
     def log(self, msg):