Skip to content

GitLab

Unverified Commit 1ebcd6c5 authored by Jeff Rasley's avatar Jeff Rasley Committed by GitHub
Browse files

Update test_sparse_attention.py

parent 1661e830
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
tests/unit/test_sparse_attention.py
View file @ 1ebcd6c5
...@@ -232,6 +232,7 @@ def init_softmax_inputs(Z, H, M, N, scale, rho, block, dtype, dense_x=True, layo ...@@ -232,6 +232,7 @@ def init_softmax_inputs(Z, H, M, N, scale, rho, block, dtype, dense_x=True, layo
def _skip_on_cuda_compatability(): def _skip_on_cuda_compatability():
pytest.skip("Skip these tests for now until we get our docker image fixed.")
if torch.cuda.get_device_capability()[0] != 7: if torch.cuda.get_device_capability()[0] != 7:
pytest.skip("needs compute capability 7; v100") pytest.skip("needs compute capability 7; v100")
cuda_major = int(torch.version.cuda.split('.')[0]) * 10 cuda_major = int(torch.version.cuda.split('.')[0]) * 10
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment