From 474fac22564f80df5b26461c61d15f2e655fb3f6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 20 Dec 2024 02:46:29 -0800 Subject: [PATCH] Update mistral.py --- unsloth/models/mistral.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index cdc73c7610..dda430459a 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -39,14 +39,14 @@ pass def MistralAttention_fast_forward( self, - hidden_states: torch.Tensor, - causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - output_attentions: bool = False, - use_cache: bool = False, - padding_mask: Optional[torch.LongTensor] = None, + hidden_states: torch.Tensor, + causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: bool = False, + use_cache: bool = False, + padding_mask: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, *args, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: