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

display the unavailable torch device on error

Browse source 
Signed-off-by: Emmanuel Ferdman <emmanuelferdman@gmail.com>
This commit is contained in:
Emmanuel Ferdman 2025-05-28 07:51:22 -07:00
parent ce75fcd7dd
commit 887930dee1
No known key found for this signature in database
GPG key ID: 0CF84597DD87A9FF
1 changed files with 2 additions and 2 deletions
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("The device: {torch_device} is not available")
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")
RuntimeError(f"The device: {torch_device} is not available")
with torch.no_grad():
stream = TextStreamer(tokenizer)