Update vision_encoder.py

vision_encoder.py

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