From 29b79d3dd2b4875ab7906b83bde61b8e9a7426e4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 24 Dec 2024 00:18:47 -0800 Subject: [PATCH] Update gemma2.py --- unsloth/models/gemma2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 4a6f3e97cc..0f0a020717 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -169,7 +169,7 @@ pass def Gemma2DecoderLayer_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,