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()