[Bugfix] Fix regression on pooling models from PR#29621 (#29921)

Signed-off-by: Roger Wang <hey@rogerw.io>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Roger Wang 2025-12-02 18:33:45 -08:00 committed by GitHub
parent 5cdd664509
commit 4dd7978374
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -134,11 +134,17 @@ class EmbeddingItems(
or a list of embedding tensors (one per item).
"""
def _unwrap(
self, item: torch.Tensor | MediaWithBytes[torch.Tensor]
) -> torch.Tensor:
"""Extract media from wrapper if present."""
return item.media if isinstance(item, MediaWithBytes) else item
def get_count(self) -> int:
return len(self.data)
def get(self, index: int) -> torch.Tensor:
return self.data[index]
return self._unwrap(self.data[index])
def get_processor_data(self) -> Mapping[str, object]:
return {}