From 03ee2e134d6cc60c4cc84096c8492fcc971f0526 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sat, 17 Oct 2026 18:10:08 -0700 Subject: [PATCH] [CI][ROCm] Fix order-dependent test_rocm_aiter_sampler_defers_import (CPU tensors hit Triton) (#62213) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 Co-authored-by: Andreas Karatzas --- tests/v1/sample/test_topk_topp_sampler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 0cd44fe213..93b71a5d95 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -110,8 +110,8 @@ def test_rocm_aiter_sampler_defers_import_when_generators_force_native( monkeypatch.setattr("builtins.__import__", guard_aiter_sampling_import) sampler = topk_topp_sampler.TopKTopPSampler() - logits = torch.randn(2, 8) - k = torch.full((2,), 2, dtype=torch.int32) + logits = torch.randn(2, 8, device=DEVICE_TYPE) + k = torch.full((2,), 2, dtype=torch.int32, device=DEVICE_TYPE) generators = {0: torch.Generator(device=logits.device).manual_seed(0)} token_ids, logits_to_return = sampler(logits, generators, k, None)