use last_hidden_state as the last layer outputs

SeanLee97 · Jan 10, 2025 · 9284508 · 9284508
1 parent aaee300
commit 9284508
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/angle_emb/angle.py b/angle_emb/angle.py
@@ -717,6 +717,7 @@ def __call__(self,
         else:
             ret = self.model(output_hidden_states=True, return_dict=True, **inputs)
             all_layer_outputs = ret.hidden_states
+            all_layer_outputs[-1] = ret.last_hidden_state
             if return_all_layer_outputs:
                 return (all_layer_outputs, ret.logits) if return_mlm_logits else all_layer_outputs
             outputs = all_layer_outputs[layer_index]