From f3a7befffa36170e0266fdc48d3d571a9bd5fa44 Mon Sep 17 00:00:00 2001
From: "Wang, Yi" <yi.a.wang@intel.com>
Date: Mon, 30 Jan 2023 22:23:55 +0800
Subject: [PATCH] fix the issue that the output dict of jit model could not get
 [0] (#21354)

---
 src/transformers/pipelines/token_classification.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/transformers/pipelines/token_classification.py b/src/transformers/pipelines/token_classification.py
index 963f491db7..e126bb7b1e 100644
--- a/src/transformers/pipelines/token_classification.py
+++ b/src/transformers/pipelines/token_classification.py
@@ -239,7 +239,8 @@ class TokenClassificationPipeline(Pipeline):
         if self.framework == "tf":
             logits = self.model(model_inputs.data)[0]
         else:
-            logits = self.model(**model_inputs)[0]
+            output = self.model(**model_inputs)
+            logits = output["logits"] if isinstance(output, dict) else output[0]
 
         return {
             "logits": logits,