[Bugfix] Flatten >2D multimodal embeddings, not just 3D (#50250)

Signed-off-by: Michal Ganczarenko <michal.ganczarenko@intel.com>
This commit is contained in:
Michał Ganczarenko
2026-08-04 13:34:58 +00:00
committed by GitHub
parent 5f2ee2fa8c
commit 7c40d61eaf
@@ -558,8 +558,8 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
total_expected = sum(split_sizes)
# Flatten to 2D: [total_tokens, hidden_dim]
if embeddings.ndim == 3:
embeddings = embeddings.view(-1, embeddings.shape[-1])
if embeddings.ndim > 2:
embeddings = embeddings.reshape(-1, embeddings.shape[-1])
total_tokens = embeddings.shape[0]
if total_tokens == total_expected: