rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-02-05 11:57:45 +01:00
Code Issues Projects Releases Wiki Activity

Fix "Expected all tensors to" on non-rotary models

Browse Source 
Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking arugment for argument index in method wrapper_index_select)
This commit is contained in:
Gnome Ann 2021-08-20 11:34:31 -04:00
parent 5f82e5ba0d
commit 56c9dc2c04
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
aiserver.py
View File

@ -386,6 +386,8 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
model.lm_head.to(breakmodel.gpu_device)
model.transformer.wte.to(breakmodel.gpu_device)
model.transformer.ln_f.to(breakmodel.gpu_device)
if(not model.config.rotary):
model.transformer.wpe.to(breakmodel.gpu_device)
gc.collect()
if(vars.bmsupported and args.breakmodel):
breakmodel.ram_blocks = max(0, min(n_layers, args.breakmodel))
@ -434,6 +436,8 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
model.lm_head.to(breakmodel.gpu_device)
model.transformer.wte.to(breakmodel.gpu_device)
model.transformer.ln_f.to(breakmodel.gpu_device)
if(not model.config.rotary):
model.transformer.wpe.to(breakmodel.gpu_device)
gc.collect()
if(vars.bmsupported and args.breakmodel):
breakmodel.ram_blocks = max(0, min(n_layers, args.breakmodel))