[FIX] better device allocation for model

NINAnor · Sep 4, 2023 · 805784b · 805784b
1 parent fa64650
commit 805784b
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/predict.py b/src/predict.py
@@ -23,10 +23,9 @@
                     format='%(asctime)s - %(levelname)s - %(message)s')
 
 def initModel(model_path, device):
-    model = torch.load(model_path)
+    model = torch.load(model_path, map_location=torch.device(device))
     model.eval()
-
-    return model.to(device)
+    return model
 
 def compute_hr(array):