From 3649ba9fa4f67ff05369d725a25e1cd4860da29d Mon Sep 17 00:00:00 2001
From: Gnome Ann <>
Date: Wed, 6 Oct 2021 12:04:56 -0400
Subject: [PATCH] Breakmodel's CUDA stream should be on primary device

---
 breakmodel.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/breakmodel.py b/breakmodel.py
index 73e40222..5724d4e2 100644
--- 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)):