make test_speculative_decoding_non_distil device-agnostic (#38010)

* make device-agnostic

* use condition

---------

Co-authored-by: Yih-Dar <2521628+ydshieh@users.noreply.github.com>
This commit is contained in:
Fanli Lin
2025-05-09 03:19:47 +08:00
committed by GitHub
parent d23aae2b8c
commit c7c2f08994

View File

@@ -2495,9 +2495,9 @@ class WhisperModelIntegrationTests(unittest.TestCase):
self.assertTrue(total_time_non_assist > total_time_assist, "Make sure that assistant decoding is faster")
@slow
@require_torch_gpu
@require_torch_accelerator
def test_speculative_decoding_non_distil(self):
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
torch_dtype = torch.float16 if torch_device in ["cuda", "xpu"] else torch.float32
model_id = "openai/whisper-large-v2"
model = WhisperForConditionalGeneration.from_pretrained(
model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True