diff --git a/src/transformers/pipelines.py b/src/transformers/pipelines.py index 5871890b..582d3fff 100755 --- a/src/transformers/pipelines.py +++ b/src/transformers/pipelines.py @@ -656,8 +656,8 @@ class TextClassificationPipeline(Pipeline): def __call__(self, *args, **kwargs): outputs = super().__call__(*args, **kwargs) - scores = np.exp(outputs) / np.exp(outputs).sum(-1) - return [{"label": self.model.config.id2label[item.argmax()], "score": item.max()} for item in scores] + scores = np.exp(outputs) / np.exp(outputs).sum(-1, keepdims=True) + return [{"label": self.model.config.id2label[item.argmax()], "score": item.max().item()} for item in scores] class FillMaskPipeline(Pipeline):