From b188bab4417f7f94df0c4e84399d163e6c0db316 Mon Sep 17 00:00:00 2001
From: rasmith <Randall.Smith@amd.com>
Date: Wed, 25 Feb 2026 13:18:00 -0600
Subject: [PATCH] [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: Randall Smith <Randall.Smith@amd.com>
---
 tests/lora/test_punica_ops.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py
index 963260367..82db7fece 100644
--- a/tests/lora/test_punica_ops.py
+++ b/tests/lora/test_punica_ops.py
@@ -395,6 +395,7 @@ def test_kernels(
     Tests LoRA kernels.
     """
     torch.set_default_device(device)
+    torch.cuda.set_device(device)
     set_random_seed(seed)
 
     if op_type == "shrink":
@@ -447,6 +448,7 @@ def test_kernels_hidden_size(
     Tests SGMV and LoRA kernels.
     """
     torch.set_default_device(device)
+    torch.cuda.set_device(device)
     set_random_seed(seed)
 
     if op_type == "shrink":