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,