Update src/omniglue/dino_extract.py

src/omniglue/dino_extract.py

@@ -28,6 +28,7 @@ class DINOExtract:
   def __init__(self, cpt_path: str, feature_layer: int = 1):
     self.feature_layer = feature_layer
     self.model = dino.vit_base()
+    self.device = "cuda" if torch.cuda.is_available() else "cpu"
     state_dict_raw = torch.load(cpt_path, map_location='cpu')
 
     # state_dict = {}
@@ -35,7 +36,7 @@ class DINOExtract:
     #   state_dict[k.replace('blocks', 'blocks.0')] = v
 
     self.model.load_state_dict(state_dict_raw)
-    self.model.eval()
+    self.model.eval().to(self.device)
 
     self.image_size_max = 630
 
@@ -98,6 +99,7 @@ class DINOExtract:
     image_processed = image / 255.0
     image_processed = (image_processed - mean) / std
     image_processed = torch.from_numpy(image_processed).permute(2, 0, 1)
+    image_processed = image_processed.to(self.device)
     return image_processed
 
   def extract_feature(self, image):