Ensure logit padding happens on default stream

PedroPareja · Aug 27, 2024 · 8d3d4c2 · 8d3d4c2
1 parent d9f0ecc
commit 8d3d4c2
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/exllamav2/model.py b/exllamav2/model.py
@@ -989,6 +989,9 @@ def forward_chunk(self,
         if self.tp_context:
             self.tp_context.wait_streams()
 
+        if x is not None and x.is_cuda:
+            torch.cuda.set_stream(torch.cuda.default_stream(x.device))
+
         # Apply logit scale
 
         # if x is not None and self.config.logit_scale != 1: