Skip to content

GitLab

"ssh:/git@developer.sourcefind.cn:2222/OpenDAS/vllm_cscc.git" did not exist on "1092a467d8562b3f61ef494ae0c33c82f0c8fff0"
Unverified Commit b188bab4 authored by rasmith's avatar rasmith Committed by GitHub
Browse files

[CI][AMD][BugFix] Add torch.cuda.set_device to test_punica_ops so punica... 


[CI][AMD][BugFix] Add  torch.cuda.set_device to test_punica_ops so punica kernels execute on same device as tensor (#34985)
Signed-off-by: default avatarRandall Smith <Randall.Smith@amd.com>
parent 15d76f74
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
tests/lora/test_punica_ops.py
View file @ b188bab4
...@@ -395,6 +395,7 @@ def test_kernels( ...@@ -395,6 +395,7 @@ def test_kernels(
Tests LoRA kernels. Tests LoRA kernels.
""" """
torch.set_default_device(device) torch.set_default_device(device)
torch.cuda.set_device(device)
set_random_seed(seed) set_random_seed(seed)
if op_type == "shrink": if op_type == "shrink":
...@@ -447,6 +448,7 @@ def test_kernels_hidden_size( ...@@ -447,6 +448,7 @@ def test_kernels_hidden_size(
Tests SGMV and LoRA kernels. Tests SGMV and LoRA kernels.
""" """
torch.set_default_device(device) torch.set_default_device(device)
torch.cuda.set_device(device)
set_random_seed(seed) set_random_seed(seed)
if op_type == "shrink": if op_type == "shrink":
......
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