to device

app.py

@@ -33,7 +33,7 @@ transform = T.Compose([
 height, width = patch_size * patch_height_nums, patch_size * patch_width_nums
 
 # DINOV2
-model = torch.hub.load(DINOV2_REPO, DINOV2_MODEL)
+model = torch.hub.load(DINOV2_REPO, DINOV2_MODEL).to(DEVICE)
 
 # faiss
 K= 5
@@ -46,7 +46,7 @@ def read_image(image_path: str) -> np.ndarray:
 
 def infer(image: np.ndarray) -> np.ndarray:
     image = np.transpose(image, (2, 0, 1)) // 255
-    transformed_image = transform(torch.Tensor(image))
+    transformed_image = transform(torch.Tensor(image)).to(DEVICE)
     embedding = model(torch.unsqueeze(transformed_image, 0))
     return embedding.detach().cpu().numpy()