Fixing single candidate_label return. (#24023)

c2e3fa0b · Nicolas Patry · GitHub · 6307312d · c2e3fa0b
Unverified Commit c2e3fa0b authored Jun 06, 2023 by Nicolas Patry Committed by GitHub Jun 06, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

src/transformers/pipelines/zero_shot_image_classification.py src/transformers/pipelines/zero_shot_image_classification.py +2 -0

No files found.
--- a/src/transformers/pipelines/zero_shot_image_classification.py
+++ b/src/transformers/pipelines/zero_shot_image_classification.py
@@ -141,6 +141,8 @@ class ZeroShotImageClassificationPipeline(Pipeline):
        if self.framework == "pt":
            probs = logits.softmax(dim=-1).squeeze(-1)
            scores = probs.tolist()
+            if not isinstance(scores, list):
+                scores = [scores]
        elif self.framework == "tf":
            probs = stable_softmax(logits, axis=-1)
            scores = probs.numpy().tolist()