diff --git a/model/model_vlm.py b/model/model_vlm.py
index 9a5f8f3..f9b10d7 100644
--- a/model/model_vlm.py
+++ b/model/model_vlm.py
@@ -74,7 +74,7 @@ class MiniMindVLM(MiniMindLM):
         if vision_tensors is not None and image_indices:
             vision_proj = self.vision_proj(vision_tensors)
             if len(vision_proj.shape) == 3:
-                vision_proj = vision_proj.unsqueeze(1)
+                vision_proj = vision_proj.unsqueeze(0)
             new_h = []
             for i in range(h.size(0)):
                 if i in image_indices: