fixed device not found

model.py

@@ -7,7 +7,7 @@ import torchvision.transforms as T
 class SegmentationModel:
     """The segmentation model."""
 
-    def __init__(self, path, device: torch.device):
+    def __init__(self, path, device='cuda' if torch.cuda.is_available() else 'cpu'):
         """
         Initializes the segmentation model, with the given path.