tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix: lin_attn_mask in wrong dtype

Browse Source
This commit is contained in:
NanoCode012
2025-02-06 15:25:33 +07:00
parent caa49a9d7d
commit ebd406af1d
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/axolotl/integrations/lolcats/linear_llama/linear_attention.py
View File

@@ -365,7 +365,7 @@ class LolcatsLinearAttention(nn.Module):
..., None ..., None
] # b, 1, k_len, 1 ] # b, 1, k_len, 1
else: else:
lin_attn_mask = attention_mask[:, None, :, None] # b, 1, k_len, 1 lin_attn_mask = attention_mask.bool()[:, None, :, None] # b, 1, k_len, 1
k = k.masked_fill(~lin_attn_mask, 0) k = k.masked_fill(~lin_attn_mask, 0)
if past_key_value is not None: # Initialize states if past_key_value is not None: # Initialize states