From 969da7d70bc0539f6be12027b71bef758325a61a Mon Sep 17 00:00:00 2001 From: Roger Wang <136131678+ywang96@users.noreply.github.com> Date: Fri, 13 Dec 2024 03:09:30 -0800 Subject: [PATCH] [V1][VLM] Fix edge case bug for InternVL2 (#11165) Signed-off-by: Roger Wang --- vllm/model_executor/models/internvl.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py index 42c769f79e202..f4b7e4478c164 100644 --- a/vllm/model_executor/models/internvl.py +++ b/vllm/model_executor/models/internvl.py @@ -669,8 +669,11 @@ def _process_image_input( image_embeds = self.extract_feature(image_input["data"]) patches_per_image = image_input["patches_per_image"] + + # Only one image in the current batch if len(patches_per_image) == 1: - image_embeds = image_embeds.unsqueeze(0) + image_embeds = image_embeds.view( + -1, self.config.text_config.hidden_size).unsqueeze(0) return image_embeds # NOTE: Image embeddings are split into separate tensors for each image