Skip to content

Commit b188bab

Browse files
authored
[CI][AMD][BugFix] Add torch.cuda.set_device to test_punica_ops so punica kernels execute on same device as tensor (vllm-project#34985)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
1 parent 15d76f7 commit b188bab

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

tests/lora/test_punica_ops.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -395,6 +395,7 @@ def test_kernels(
395395
Tests LoRA kernels.
396396
"""
397397
torch.set_default_device(device)
398+
torch.cuda.set_device(device)
398399
set_random_seed(seed)
399400

400401
if op_type == "shrink":
@@ -447,6 +448,7 @@ def test_kernels_hidden_size(
447448
Tests SGMV and LoRA kernels.
448449
"""
449450
torch.set_default_device(device)
451+
torch.cuda.set_device(device)
450452
set_random_seed(seed)
451453

452454
if op_type == "shrink":

0 commit comments

Comments
 (0)