Update inference.py

inference.py

@@ -20,7 +20,7 @@ class CoralSegModel:
         rng = np.random.RandomState(0)
         self.palette = (rng.randint(0, 255, size=(num_classes, 3))).astype(np.uint8)
 
-    @torch.inference_mode()
+    @spaces.GPU
     def predict_overlay(self, frame_bgr: np.ndarray, alpha: float = 0.45) -> np.ndarray:
         """
         frame_bgr: np.ndarray HxWx3 in BGR (as read by OpenCV)
@@ -30,7 +30,7 @@ class CoralSegModel:
         rgb = frame_bgr[:, :, ::-1]
         pil = Image.fromarray(rgb)
 
-        inputs = self.processor(images=pil, return_tensors="pt", device=self.device)
+        inputs = self.processor(images=pil, return_tensors="pt", device=self.device).to(device)
         outputs = self.model(**inputs)
         logits = outputs.logits  # [B, C, h, w]
         upsampled = torch.nn.functional.interpolate(