Merge pull request #115 from CgggggY/fix-single-batch-multi-image-shape

修复 bs=1 多图输入时 vision_tensors 维度不一致的问题

model/model_vlm.py

@@ -115,12 +115,18 @@ class MiniMindVLM(MiniMindForCausalLM):
             if hasattr(pixel_values, 'keys'):
                 img_emb = MiniMindVLM.get_image_embeddings(pixel_values, self.vision_encoder)
                 vision_tensors = self.vision_proj(img_emb)
+
             else:
                 if len(pixel_values.shape) == 6:
                     pixel_values = pixel_values.squeeze(2)
                 bs, num, c, im_h, im_w = pixel_values.shape
-                stack_dim = 1 if bs > 1 else 0
-                vision_tensors = torch.stack([self.vision_proj(MiniMindVLM.get_image_embeddings(pixel_values[:, i, :, :, :], self.vision_encoder)) for i in range(num)], dim=stack_dim)
+
+                # 每次取整个 batch 的第 i 张图，经过 vision encoder + projector 后得到 [B, target_tokens, hidden_size]
+                # 再统一在图片维（dim=1）上堆叠，使 vision_tensors 始终为 [B, N, target_tokens, hidden_size]
+                # 这样可与后续 count_vision_proj 中的 vf[b][k]、vf.size(1) 保持一致
+                vision_tensors = torch.stack(
+                    [self.vision_proj(MiniMindVLM.get_image_embeddings(pixel_values[:, i, :, :, :], self.vision_encoder))
+                        for i in range(num)],dim=1)
             hidden_states = self.count_vision_proj(tokens=input_ids, h=hidden_states, vision_tensors=vision_tensors, seqlen=input_ids.shape[1])
 
         position_embeddings = (
@@ -153,4 +159,4 @@ class MiniMindVLM(MiniMindForCausalLM):
             loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100)
 
         output = MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=presents, hidden_states=hidden_states)
-        return output
+        return output