fix the issue that the output dict of jit model could not get [0] (#21354)
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user