From c706447fd4bbace502e6ab2de214d3de38e9961d Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sat, 24 Feb 2024 20:04:31 +1100 Subject: [PATCH] Update gemma.py --- unsloth/models/gemma.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 1d0960f5f0..114d98e95f 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -180,9 +180,9 @@ def GemmaDecoderLayer_fast_forward( past_key_value: Optional[Tuple[torch.Tensor]] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, - padding_mask: Optional[torch.LongTensor] = None, + cache_position: Optional[torch.LongTensor] = None, *args, **kwargs, -) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: +): if False:#past_key_value is not None: do_prefill = not hasattr(self.self_attn, "paged_attention") @@ -215,6 +215,7 @@ def GemmaDecoderLayer_fast_forward( past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, + cache_position=cache_position, # padding_mask=padding_mask, ) hidden_states = residual + hidden_states