Fixed bug key error "last_hidden_state" (#1674)

* bug fix: key error last_hidden_state * fix --------- Co-authored-by: Satish Silveri <[email protected]> Co-authored-by: fxmarty <[email protected]>
huggingface · Jul 1, 2024 · a5500c7 · a5500c7
1 parent d82d4c6
commit a5500c7
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py
@@ -1091,10 +1091,10 @@ def forward(
             onnx_outputs = self.model.run(None, onnx_inputs)
             model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs)
 
-            # TODO: why do we only return last_hidden_state? why not all outputs?
-            # that way, there will be less need for ORTModelForCustomTask in cases where
-            # we just want to extend model outputs with attentions, hidden_states, etc.
-            last_hidden_state = model_outputs["last_hidden_state"]
+        if "last_hidden_state" in self.output_names:
+            last_hidden_state = model_outputs[self.output_names["last_hidden_state"]]
+        else:
+            last_hidden_state = model_outputs[0]
 
         # converts output to namedtuple for pipelines post-processing
         return BaseModelOutput(last_hidden_state=last_hidden_state)