rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Breakmodel's CUDA stream should be on primary device

Browse Source
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
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
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)):