From 1c4a2c9a71b528de396c51c1375975c9d88b72f2 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Fri, 18 Apr 2025 06:17:58 -0700
Subject: [PATCH] Make exllamav3 safer as well

---
 modules/exllamav3_hf.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/modules/exllamav3_hf.py b/modules/exllamav3_hf.py
index 3bf44c9b..0f742fa2 100644
--- a/modules/exllamav3_hf.py
+++ b/modules/exllamav3_hf.py
@@ -153,6 +153,9 @@ class Exllamav3HF(PreTrainedModel):
         else:
             self.past_seq = seq_tensor
 
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
         loss = None
         if labels is not None:
             # Shift so that tokens < n predict n