Update seg.py

seg.py

@@ -6,7 +6,8 @@ import cv2
 
 # 아래는 unet이고 dinov2 전에 실험했던 모델입니다. 
 def segmentation(img):
-    device=torch.device("cuda")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    #device=torch.device("cuda")
     net=torch.load("weights/unet.pth")
     
     mask_values=[[0, 0, 0], [255, 255, 255]]