Update vision_encoder.py

vision_encoder.py

@@ -1,10 +1,10 @@
-import torch.nn as nn
-from transformers import ViTModel
+from transformers import  ViTModel
 from torchvision import transforms
-import torch
+import torch 
 
 import transformers
 
+
 transformers.logging.set_verbosity_error()
 
 class VisionEncoder(nn.Module):
@@ -17,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, image, device):
-        processed_images = torch.stack([self.image_transform(image) for image in images]).to(device)        
+    def forward(self, images,device):
+        processed_images = torch.stack([self.image_transform(image) for image in images]).to(device)
         with torch.no_grad():
-            pixel_values = self.vision_model(processed_image)
+            pixel_values = self.vision_model(processed_images)
             image_features = pixel_values.last_hidden_state
         return image_features