From d8bc6402b59bd5f349b06e2ed1e6056a5dafa569 Mon Sep 17 00:00:00 2001 From: Emmanuel Ferdman Date: Thu, 29 May 2025 10:37:41 +0300 Subject: [PATCH] raise exception on device error (#1342) * display the unavailable torch device on error * Raise exception on device error --------- Signed-off-by: Emmanuel Ferdman --- ktransformers/util/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ktransformers/util/utils.py b/ktransformers/util/utils.py index 7301572..308def1 100644 --- a/ktransformers/util/utils.py +++ b/ktransformers/util/utils.py @@ -243,7 +243,7 @@ def prefill_and_generate(model, tokenizer, inputs, max_new_tokens=10000, use_cud elif torch.xpu.is_available(): torch.xpu.set_device(torch_device) else: - RuntimeError("The device: {torch_device} is not available") + raise RuntimeError(f"The device: {torch_device} is not available") inputs_embeds = model.model.embed_tokens(cur_token.to("cpu")).to(torch_device) # with torch.cuda.stream(custom_stream): logits=model(inputs_embeds=inputs_embeds, @@ -284,7 +284,7 @@ def prefill_and_generate(model, tokenizer, inputs, max_new_tokens=10000, use_cud elif torch.xpu.is_available(): torch.xpu.set_device(torch_device) else: - RuntimeError("The device: {torch_device} is not available") + raise RuntimeError(f"The device: {torch_device} is not available") with torch.no_grad(): stream = TextStreamer(tokenizer)