diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py
index 2c43cd72147..f517abb2b14 100644
--- a/tests/v1/worker/test_gpu_input_batch_v2.py
+++ b/tests/v1/worker/test_gpu_input_batch_v2.py
@@ -5,9 +5,10 @@
 import pytest
 import torch
 
+from vllm.platforms import current_platform
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 
-DEVICE = "cuda"
+DEVICE = current_platform.device_type
 
 
 @pytest.mark.parametrize(