From b0f97ccb34b4f4f0b11d0bdca8ddbe6b3e5d5b8e Mon Sep 17 00:00:00 2001 From: Emmanuel Ferdman Date: Thu, 29 May 2025 00:25:25 -0700 Subject: [PATCH] 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 4864a4c..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(f"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(f"The device: {torch_device} is not available") + raise RuntimeError(f"The device: {torch_device} is not available") with torch.no_grad(): stream = TextStreamer(tokenizer)