Elgato_dark/unsloth
1
0
Fork 0
mirror of https://github.com/unslothai/unsloth synced 2026-04-21 13:37:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

attention mask

Browse source
This commit is contained in:
Daniel Han-Chen 2024-01-28 04:18:19 +11:00
parent 6c7f0dbcb4
commit 166f8c812e
2 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
unsloth/models/llama.py
View file

@ -489,7 +489,7 @@ def LlamaModel_fast_forward(
# Ignore attention_mask
if attention_mask is None:
padding_mask = None
elif False:#self.training:
elif self.training:
attention_mask = None
padding_mask = None
else:

2
unsloth/models/mistral.py
View file

@ -90,7 +90,7 @@ def MistralAttention_fast_forward(
past_key_value = (K, V) if use_cache else None
# Attention module
if (attention_mask is None and not HAS_FLASH_ATTENTION):
if (not HAS_FLASH_ATTENTION):
# Xformers memory efficient attention
Q = Q.transpose(1, 2)
K = K.transpose(1, 2)