Update gemma.py

This commit is contained in:
Daniel Han-Chen 2024-02-24 20:04:31 +11:00
commit c706447fd4

View file

@ -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