vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2025-09-08 21:49:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Raise exception on device error

Browse source 
Signed-off-by: Emmanuel Ferdman <emmanuelferdman@gmail.com>
This commit is contained in:
Emmanuel Ferdman 2025-05-29 00:25:25 -07:00
parent 887930dee1
commit b0f97ccb34
No known key found for this signature in database
GPG key ID: 0CF84597DD87A9FF
1 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
ktransformers/util/utils.py
View file

@ -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)