Update model_utils.py

model_utils.py

@@ -2,7 +2,7 @@ import torch
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
 import whisper
 from config import WHISPER_MODEL_SIZE
-import spaces 
+import spaces
 
 # Global variables to store models
 whisper_processor = None
@@ -15,13 +15,13 @@ def load_models():
     if whisper_processor is None:
         whisper_processor = WhisperProcessor.from_pretrained(f"openai/whisper-{WHISPER_MODEL_SIZE}")
     if whisper_model is None:
-        whisper_model = WhisperForConditionalGeneration.from_pretrained(f"openai/whisper-{WHISPER_MODEL_SIZE}").to(get_device())
+        whisper_model = WhisperForConditionalGeneration.from_pretrained(f"openai/whisper-{WHISPER_MODEL_SIZE}")
     if whisper_model_small is None:
         whisper_model_small = whisper.load_model(WHISPER_MODEL_SIZE)
 
 @spaces.GPU
 def get_device():
-    return "cuda:0" if torch.cuda.is_available() else "cpu"
+    return "cuda" if torch.cuda.is_available() else "cpu"
 
 @spaces.GPU
 def get_processor():
@@ -35,7 +35,7 @@ def get_model():
     global whisper_model
     if whisper_model is None:
         load_models()
-    return whisper_model
+    return whisper_model.to(get_device())
 
 @spaces.GPU
 def get_whisper_model_small():