Skip to content

GitLab

Commit ad50bf4b authored by simon-mo's avatar simon-mo
Browse files

fix lint

parent cf6ff182
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
tests/kernels/test_prefix_prefill.py
View file @ ad50bf4b
......@@ -36,8 +36,8 @@ def test_contexted_kv_attention(
torch.cuda.manual_seed(0)
torch.set_default_device(device)
# Need this, otherwise when we capture the graph the process for GPU 1 would run on both
# GPU0 and GPU1 and things would hang
# Need this, otherwise when we capture the graph the process for GPU 1 would
# run on both GPU0 and GPU1 and things would hang
#
# see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523
torch.cuda.set_device(device)
......
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