reverted to cuda check for gpu

joshsoftware · Aug 23, 2024 · 6b2c44e · 6b2c44e
1 parent 2a245d5
commit 6b2c44e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/load_model.py b/load_model.py
@@ -4,8 +4,8 @@
 #load the whisper model from net if it isn't stored locally
 def load_model(model_id, model_path):
     #check GPU is avaialbe
-    #device = "cuda" if torch.cuda.is_available() else "cpu"
-    device = "cpu"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    #device = "cpu"
     model = whisper.load_model(model_id, device=device, download_root=model_path)
     print(
         f"Model will be run on {device}\n"