Upload model

modeling.py

@@ -405,8 +405,14 @@ class TotalClassifierModel(PreTrainedModel):
         return slices
 
     def preprocess(
-        self, x: np.ndarray, mode: str = "2d", torchify: bool = True
+        self,
+        x: np.ndarray,
+        mode: str = "2d",
+        torchify: bool = True,
+        device: str | torch.device | None = None,
     ) -> np.ndarray:
+        if device is not None:
+            assert torchify, "`torchify` must be `True` if specifying `device`"
         mode = mode.lower()
         if mode == "2d":
             x = cv2.resize(x, self.image_size)
@@ -421,6 +427,8 @@ class TotalClassifierModel(PreTrainedModel):
                 x = rearrange(torch.from_numpy(x).float(), "h w c -> c h w")
             elif x.ndim == 4:
                 x = rearrange(torch.from_numpy(x).float(), "n h w c -> n c h w")
+        if device is not None:
+            x = x.to(device)
         return x
 
     def crop_single_plane(