Update image_to_text.py

image_to_text.py

@@ -17,7 +17,8 @@ def generate_caption(image):
     else:
         image = Image.open(image).convert('RGB')
     
-    pixel_values = feature_extractor(images=[image], return_tensors="pt").pixel_values
+    # Add padding=True to handle images of different sizes
+    pixel_values = feature_extractor(images=[image], return_tensors="pt", padding=True).pixel_values
     pixel_values = pixel_values.to(device)
 
     # Generate caption