From fa4ca3f28a7a4e5fda522a9ee8e98a54ae9b7415 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 24 Dec 2024 00:09:23 -0800 Subject: [PATCH] Update granite.py --- unsloth/models/granite.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index 2229636e9e..9466a8d6c1 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -60,7 +60,7 @@ pass def GraniteAttention_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, @@ -171,7 +171,7 @@ pass def GraniteDecoderLayer_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,