"vscode:/vscode.git/clone" did not exist on "afc83aa216249e813d41eb2b72676d0d74622004"
Unverified Commit 1ebcd6c5 authored by Jeff Rasley's avatar Jeff Rasley Committed by GitHub
Browse files

Update test_sparse_attention.py

parent 1661e830
......@@ -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():
pytest.skip("Skip these tests for now until we get our docker image fixed.")
if torch.cuda.get_device_capability()[0] != 7:
pytest.skip("needs compute capability 7; v100")
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