From af26f3c3690082dc58b37a5c2fb709abcf1f930b Mon Sep 17 00:00:00 2001 From: Kareem Horstink Date: Sun, 23 Aug 2026 18:15:50 +0000 Subject: [PATCH] fix: extract pooler_output from BaseModelOutputWithPooling in get_image_features (newer transformers) --- src/scorers/laion.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/scorers/laion.py b/src/scorers/laion.py index 38b9bce..5e55ffb 100644 --- a/src/scorers/laion.py +++ b/src/scorers/laion.py @@ -131,7 +131,10 @@ class LaionScorer: inputs = {k: v.to(self._device) for k, v in inputs.items()} 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) score = self.head(embeds).squeeze().item()