fix: extract pooler_output from BaseModelOutputWithPooling in get_image_features (newer transformers)

This commit is contained in:
Kareem Horstink
2026-08-23 18:15:50 +00:00
parent 016a68e0a2
commit af26f3c369
+4 -1
View File
@@ -131,7 +131,10 @@ class LaionScorer:
inputs = {k: v.to(self._device) for k, v in inputs.items()} inputs = {k: v.to(self._device) for k, v in inputs.items()}
with torch.no_grad(): with torch.no_grad():
embeds = self.clip_model.get_image_features(**inputs) out = self.clip_model.get_image_features(**inputs)
# Newer transformers return a `BaseModelOutputWithPooling`
# object; older ones return the tensor directly. Handle both.
embeds = getattr(out, "pooler_output", out)
embeds = nn.functional.normalize(embeds, dim=-1) embeds = nn.functional.normalize(embeds, dim=-1)
score = self.head(embeds).squeeze().item() score = self.head(embeds).squeeze().item()