Fix zero-shot pipeline single seq output shape (#6104)

huggingface · Jul 28, 2020 · b1c8b76 · b1c8b76
1 parent 06834bc
commit b1c8b76
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/pipelines.py b/src/transformers/pipelines.py
@@ -959,7 +959,7 @@ def __call__(self, sequences, candidate_labels, hypothesis_template="This exampl
             top_inds = list(reversed(scores[iseq].argsort()))
             result.append(
                 {
-                    "sequence": sequences if num_sequences == 1 else sequences[iseq],
+                    "sequence": sequences if isinstance(sequences, str) else sequences[iseq],
                     "labels": [candidate_labels[i] for i in top_inds],
                     "scores": scores[iseq][top_inds].tolist(),
                 }