Upload model

modeling.py

@@ -343,6 +343,7 @@ class CTCropModel(PreTrainedModel):
         rescaled_coords[:, 3] = rescaled_coords[:, 3] * img_shape[:, 0]
         return rescaled_coords.int()
 
+    @torch.no_grad()
     def crop(
         self,
         x: np.ndarray,
@@ -351,7 +352,7 @@ class CTCropModel(PreTrainedModel):
         raw_hu: bool = False,
         remove_empty_slices: bool = False,
         add_buffer: float | tuple[float, float] | None = None,
-    ) -> np.ndarray:
+    ) -> np.ndarray | tuple[np.ndarray, list[int]]:
         assert mode in ["2d", "3d"]
         if device is None:
             device = "cuda" if torch.cuda.is_available() else "cpu"