Skip to content

Commit 01dacb9

Browse files
committed
missing tqdm dep for tests
1 parent 1b367c6 commit 01dacb9

File tree

2 files changed

+10
-2
lines changed

2 files changed

+10
-2
lines changed

pyproject.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,8 @@ examples = [
4242
"wandb"
4343
]
4444
test = [
45-
"pytest"
45+
"pytest",
46+
"tqdm",
4647
]
4748

4849
[tool.pytest.ini_options]

tests/test_sparse_attn.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,14 @@ def test_transformer_inference():
9090
num_tokens = 256,
9191
dim = 512,
9292
depth = 2,
93-
use_sparse_attn = True
93+
causal = True,
94+
use_sparse_attn = True,
95+
sparse_attn_kwargs = dict(
96+
sliding_window_size = 16,
97+
compress_block_size = 4,
98+
selection_block_size = 16,
99+
num_selected_blocks = 1
100+
)
94101
)
95102

96103
prompt = torch.randint(0, 256, (1, 1))

0 commit comments

Comments
 (0)