Breakmodel's CUDA stream should be on primary device

2021-10-06 12:04:56 -04:00 · 2021-10-06 12:04:56 -04:00 · 3649ba9fa4
parent aa59f8b4b2
commit 3649ba9fa4
1 changed files with 1 additions and 1 deletions
--- a/breakmodel.py
+++ b/breakmodel.py
@ -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)):