update gpu

app.py

@@ -40,8 +40,6 @@ class MobileSAM(nn.Module):
             with open(sam_checkpoint, 'wb') as f:
                 f.write(r.content)
 
-        device = 'cuda' if USE_CUDA else 'cpu'
-
         mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
 
         def new_forward_fn(self, x):
@@ -123,7 +121,6 @@ class MobileSAM(nn.Module):
         setattr(mobile_sam.image_encoder.layers[1].blocks[0].__class__, "forward", new_forward_fn2)
 
 
-        mobile_sam.to(device=device)
         mobile_sam.eval()
         self.image_encoder = mobile_sam.image_encoder
         
@@ -150,7 +147,8 @@ def image_mobilesam_feature(
     images,
     node_type="block",
     layer=-1,
-):
+):  
+    print("Running MobileSAM")
     global USE_CUDA
     if USE_CUDA:
         images = images.cuda()
@@ -160,6 +158,7 @@ def image_mobilesam_feature(
     if USE_CUDA:
         feat_extractor = feat_extractor.cuda()
 
+    print("images shape:", images.shape)
     # attn_outputs, mlp_outputs, block_outputs = [], [], []
     outputs = []
     for i in range(images.shape[0]):
@@ -228,8 +227,6 @@ class SAM(torch.nn.Module):
 
         self.image_encoder = sam.image_encoder
         self.image_encoder.eval()
-        if USE_CUDA:
-            self.image_encoder = self.image_encoder.cuda()
 
     @torch.no_grad()
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -291,8 +288,6 @@ class DiNOv2(torch.nn.Module):
         self.dinov2 = torch.hub.load("facebookresearch/dinov2", ver)
         self.dinov2.requires_grad_(False)
         self.dinov2.eval()
-        if USE_CUDA:
-            self.dinov2 = self.dinov2.cuda()
 
         def new_block_forward(self, x: torch.Tensor) -> torch.Tensor:
             def attn_residual_func(x):
@@ -372,8 +367,6 @@ class CLIP(torch.nn.Module):
         model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
         # processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
         self.model = model.eval()
-        if USE_CUDA:
-            self.model = self.model.cuda()
 
         def new_forward(
             self,