Fix forward bug in ORTModelForFeatureExtraction (#1941)

Fix forward bug for ORTModelForFeatureExtraction
huggingface · Jul 4, 2024 · 5c803db · 5c803db
1 parent f755a58
commit 5c803db
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py
@@ -1091,11 +1091,11 @@ def forward(
             onnx_outputs = self.model.run(None, onnx_inputs)
             model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs)
 
-        if "last_hidden_state" in self.output_names:
-            last_hidden_state = model_outputs["last_hidden_state"]
-        else:
-            # TODO: This allows to support sentence-transformers models (sentence embedding), but is not validated.
-            last_hidden_state = next(iter(model_outputs.values()))
+            if "last_hidden_state" in self.output_names:
+                last_hidden_state = model_outputs["last_hidden_state"]
+            else:
+                # TODO: This allows to support sentence-transformers models (sentence embedding), but is not validated.
+                last_hidden_state = next(iter(model_outputs.values()))
 
         # converts output to namedtuple for pipelines post-processing
         return BaseModelOutput(last_hidden_state=last_hidden_state)