Update inference_seg.py

inference_seg.py

@@ -18,13 +18,16 @@ def load_model(use_box=False):
     )
     MODEL.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
     MODEL.eval()
-    DEVICE = torch.device("cpu")
+    if torch.cuda.is_available():
+        DEVICE = torch.device("cuda")
+    else:
+        DEVICE = torch.device("cpu")
     return MODEL, DEVICE
 
 
 @torch.no_grad()
 def run(model, img_path, box=None, device="cpu"):
-    print(device)
+    print("DEVICE": device)
     model.move_to_device(device)
     model.eval()
     with torch.no_grad():