Unverified Commit f3a7beff authored by Wang, Yi's avatar Wang, Yi Committed by GitHub
Browse files

fix the issue that the output dict of jit model could not get [0] (#21354)

parent c749bd40
...@@ -239,7 +239,8 @@ class TokenClassificationPipeline(Pipeline): ...@@ -239,7 +239,8 @@ class TokenClassificationPipeline(Pipeline):
if self.framework == "tf": if self.framework == "tf":
logits = self.model(model_inputs.data)[0] logits = self.model(model_inputs.data)[0]
else: else:
logits = self.model(**model_inputs)[0] output = self.model(**model_inputs)
logits = output["logits"] if isinstance(output, dict) else output[0]
return { return {
"logits": logits, "logits": logits,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment