rixty
/
KoboldAI-Client
mirror of https://github.com/KoboldAI/KoboldAI-Client.git
1
1
Fork 0
Code Issues Projects Releases Wiki Activity

Breakmodel's CUDA stream should be on primary device

This commit is contained in:
Gnome Ann 2021-10-06 12:04:56 -04:00
parent aa59f8b4b2
commit 3649ba9fa4
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
breakmodel.py
View File

@ -387,7 +387,7 @@ def new_forward(
all_hidden_states = () if output_hidden_states else None all_hidden_states = () if output_hidden_states else None
if breakmodel and ram_blocks: if breakmodel and ram_blocks:
copystream = torch.cuda.Stream(device=0,priority = -1) copystream = torch.cuda.Stream(device=primary_device, priority=-1)
for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)): for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):