[fix] Prevent IndexError if output_hidden_states & ONNX (#3008)

UKPLab · Oct 21, 2024 · f286d9f · f286d9f
1 parent a028b58
commit f286d9f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sentence_transformers/models/Transformer.py b/sentence_transformers/models/Transformer.py
@@ -352,8 +352,8 @@ def forward(self, features: dict[str, torch.Tensor], **kwargs) -> dict[str, torc
 
         features.update({"token_embeddings": output_tokens, "attention_mask": features["attention_mask"]})
 
-        if self.auto_model.config.output_hidden_states:
-            all_layer_idx = 2
+        if self.auto_model.config.output_hidden_states and len(output_states) > 2:
+            all_layer_idx = 2  # I.e. after last_hidden_states and pooler_output
             if len(output_states) < 3:  # Some models only output last_hidden_states and all_hidden_states
                 all_layer_idx = 1