Update vision_encoder.py

vision_encoder.py

@@ -1,11 +1,10 @@
-import torch.nn as nn 
-from transformers import  ViTModel
+import torch.nn as nn
+from transformers import ViTModel
 from torchvision import transforms
-import torch 
+import torch
 
 import transformers
 
-
 transformers.logging.set_verbosity_error()
 
 class VisionEncoder(nn.Module):
@@ -18,9 +17,9 @@ class VisionEncoder(nn.Module):
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
         ])
 
-    def forward(self, images,device):
-        processed_images = torch.stack([self.image_transform(image) for image in images]).to(device)
+    def forward(self, image, device):
+        processed_image = self.image_transform(image).unsqueeze(0).to(device)
         with torch.no_grad():
-            pixel_values = self.vision_model(processed_images)
+            pixel_values = self.vision_model(processed_image)
             image_features = pixel_values.last_hidden_state
-        return image_features
+        return image_features