Fix Pixtral-HF

Signed-off-by: DarkLight1337 <[email protected]>
vllm-project · Dec 21, 2024 · 5fcb5d6 · 5fcb5d6
1 parent 38aaff8
commit 5fcb5d6
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/vllm/multimodal/processing.py b/vllm/multimodal/processing.py
@@ -706,7 +706,11 @@ def _cached_call_fine(
 
                 for k, v in processed_modal_item.items():
                     # Remove the extra batch dimension (if it exists)
-                    processed_modal_items[k].append(v.squeeze(0))
+                    # NOTE: v may be a list instead of a tensor
+                    if len(v) == 1:
+                        v = v[0]
+
+                    processed_modal_items[k].append(v)
 
             for k, vs in processed_modal_items.items():
                 # Try to merge elements into a single tensor