CPU

app.py

@@ -51,7 +51,7 @@ def predict(img_input):
 
     t = T.Compose([T.ToTensor(), NORMALIZE])
     img_t = t(img)[None,:,:,:]
-    inputs = img_t.to("cuda")
+    inputs = img_t
     
     # Forward step
     with torch.no_grad():

model.py

@@ -161,7 +161,6 @@ class FoundModel(nn.Module):
         # Decoder
         self.decoder.load_state_dict(state_dict["decoder"])
         self.decoder.eval()
-        self.decoder.to("cuda")
 
 
     @torch.no_grad()
@@ -222,7 +221,7 @@ def get_vit_encoder(vit_arch, vit_model, vit_patch_size, enc_type_feats):
         # TODO change if want to have last layer not unfrozen
         for p in vit_encoder.parameters():
             p.requires_grad = False
-        vit_encoder.eval().cuda()  # mode eval
+        vit_encoder.eval() # mode eval
         state_dict = torch.hub.load_state_dict_from_url(
             url="https://dl.fbaipublicfiles.com/dino/" + url
         )