loads weights to cpu if cuda not available

model.py

@@ -15,7 +15,7 @@ class SegmentationModel:
             path: The path to the model file.
             device: The device to run the model on.
         """
-        self.model = torch.load(path)
+        self.model = torch.load(path, map_location=device)
         self.model.to(device)
         self.model.eval()
         self.device = device