Skip to content

Commit

Permalink
Fix typo in readme (#64)
Browse files Browse the repository at this point in the history
stack-info: PR: #64, branch: drisspg/stack/1
  • Loading branch information
drisspg authored Oct 23, 2024
1 parent 5475578 commit 36f8bd5
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ cd attention-gym
pip install .
```

## 💻 Usage
## 💻 Usage

There are two main ways to use Attention Gym:

Expand All @@ -48,10 +48,10 @@ There are two main ways to use Attention Gym:
```python
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
from attn_gym.masks import generate_sliding_window
# Use the imported function in your code
sliding_window_mask = generate_sliding_window(window_size=1024)
block_mask = create_block_mask(mask_mod, 1, 1, S, S, device=device)
sliding_window_mask_mod = generate_sliding_window(window_size=1024)
block_mask = create_block_mask(sliding_window_mask_mod, 1, 1, S, S, device=device)
out = flex_attention(query, key, value, block_mask=block_mask)
```

Expand Down

0 comments on commit 36f8bd5

Please sign in to comment.