Update llama.py

2024-02-04 02:42:30 +11:00 · 2024-02-04 02:42:30 +11:00 · ac9bc79251
commit ac9bc79251
parent 8c8685eeef
1 changed files with 1 additions and 0 deletions
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@ -309,6 +309,7 @@ def LlamaAttention_fast_forward(
        V = V.transpose(1, 2)
        A = flash_attn_func(Q, K, V, causal = True)
    else:
+        print(attention_mask)
        # Grouped query attention
        if n_groups != 1:
            K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)