Skip to content

Commit 4a74d6a

Browse files
committed
Fix README examples: create attn_bias and pass attention_mask=None to create_mask (EN and ZH)
1 parent b9260c0 commit 4a74d6a

File tree

2 files changed

+3
-3
lines changed

2 files changed

+3
-3
lines changed

README.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -169,13 +169,13 @@ key = torch.randn(batch_size, seq_len, num_kv_heads, head_dim, device=device, dt
169169
value = torch.randn(batch_size, seq_len, num_kv_heads, head_dim, device=device, dtype=dtype)
170170

171171
# Create bias for sparse attention
172-
attn_mask = torch.ones(batch_size, num_kv_heads, seq_len, seq_len, device=device, dtype=dtype)
172+
attn_bias = torch.ones(batch_size, num_kv_heads, seq_len, seq_len, device=device, dtype=dtype)
173173

174174
# Generate dynamic mask based on bias
175175
if seq_len > window_size:
176176
attn_mask = create_mask(
177177
attention_bias=attn_bias,
178-
attention_mask=attn_mask,
178+
attention_mask=None,
179179
batch_size=batch_size,
180180
query_len=seq_len,
181181
key_len=seq_len,

README_zh.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,7 @@ attn_bias = torch.randn(batch_size, num_kv_heads, seq_len, seq_len, device=devic
175175
if seq_len > window_size:
176176
attn_mask = create_mask(
177177
attention_bias=attn_bias,
178-
attention_mask=attn_mask,
178+
attention_mask=None,
179179
batch_size=batch_size,
180180
query_len=seq_len,
181181
key_len=seq_len,

0 commit comments

Comments
 (0)