From b3d7e6160840b23e1c4334195abe0adf52f0ccb9 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sat, 24 Feb 2024 20:20:37 +1100 Subject: [PATCH] Update gemma.py --- unsloth/models/gemma.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 02025cf1ac..243c9d9d7b 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -173,15 +173,14 @@ pass # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L590 def GemmaDecoderLayer_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: Optional[bool] = False, - use_cache: Optional[bool] = False, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, - *args, **kwargs, + **kwargs, ): if False:#past_key_value is not None: do_prefill = not hasattr(self.self_attn, "paged_attention")