Update inference_seg.py

inference_seg.py

@@ -25,7 +25,7 @@ def load_model(use_box=False):
 @torch.no_grad()
 def run(model, img_path, box=None, device="cpu"):
     output = model(img_path, box=box)
-    mask = output["pred"]
+    mask = output
     mask = (mask > 0).astype(np.uint8)
     return mask
 # import os