Move inputs to model's device

src/rad_dino/__init__.py

@@ -6,7 +6,7 @@ from torch import Tensor
 from torch import nn
 from transformers import AutoImageProcessor
 from transformers import AutoModel
-from transformers.image_processing_base import BatchFeature
+from transformers.feature_extraction_utils import BatchFeature
 
 
 __version__ = "0.1.0"
@@ -25,6 +25,10 @@ class RadDino(nn.Module):
         self.model = AutoModel.from_pretrained(self._REPO).eval()
         self.processor = AutoImageProcessor.from_pretrained(self._REPO, use_fast=False)
 
+    @property
+    def device(self) -> torch.device:
+        return next(self.model.parameters()).device
+
     def preprocess(self, image_or_images: TypeInputImages) -> BatchFeature:
         return self.processor(image_or_images, return_tensors="pt")
 
@@ -53,7 +57,7 @@ class RadDino(nn.Module):
         self,
         image_or_images: TypeInputImages,
     ) -> tuple[TypeClsToken, TypePatchTokens]:
-        inputs = self.preprocess(image_or_images)
+        inputs = self.preprocess(image_or_images).to(self.device)
         cls_token, patch_tokens_flat = self.encode(inputs)
         patch_tokens = self.reshape_patch_tokens(patch_tokens_flat)
         return cls_token, patch_tokens