Update seg.py

seg.py

@@ -8,7 +8,8 @@ import cv2
 # dinov2 모델을 segmentation 모델로 활용하였고 이를 train 하여서 pretrained weight를 weight 폴더에 저장하고 아래와 같이 load합니다.
 
 def segmentation(img):
-    device=torch.device("cuda")
+    #device=torch.device("cuda")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     net=torch.load("weights/dinov2_model.pth")
     
     mask_values=[[0, 0, 0], [255, 255, 255]]