ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update gemma.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-02-24 16:57:16 +11:00
commit 4c7d21e41b
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
unsloth/models/gemma.py
View file

@ -61,7 +61,7 @@ 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,
# padding_mask: Optional[torch.LongTensor] = None,
*args, **kwargs,
) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
if past_key_value is not None:
@ -313,7 +313,7 @@ def GemmaModel_fast_forward(
def create_custom_forward(module):
def custom_forward(*inputs):
# None for past_key_value
return module(*inputs, past_key_value, output_attentions, padding_mask=padding_mask)
return module(*inputs, past_key_value, output_attentions)#, padding_mask=padding_mask)
return custom_forward
@ -335,7 +335,7 @@ def GemmaModel_fast_forward(
past_key_value=past_key_value,
output_attentions=output_attentions,
use_cache=use_cache,
padding_mask=padding_mask,
# padding_mask=padding_mask,
)
hidden_states = layer_outputs[0]