From 00af89272cdbdc4917c3d470b57d18a6d1bc4296 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 24 Dec 2024 00:04:36 -0800 Subject: [PATCH] Update mistral.py --- unsloth/models/mistral.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index dda430459a..d6c6946664 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -40,7 +40,7 @@ pass def MistralAttention_fast_forward( self, hidden_states: torch.Tensor, - causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, + causal_mask: Optional[BlockDiagonalCausalMask] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Tuple[torch.Tensor]] = None, @@ -172,7 +172,7 @@ pass def MistralForCausalLM_fast_forward( self, input_ids: torch.LongTensor = None, - causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, + causal_mask: Optional[BlockDiagonalCausalMask] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None,