From 25e46d8b305d5dd43bdb4ff59ee3d4e9b5cd26f5 Mon Sep 17 00:00:00 2001
From: turboderp <11859846+turboderp@users.noreply.github.com>
Date: Thu, 11 Jul 2024 01:08:00 +0200
Subject: [PATCH] Dynamic gen: Fix return_logits=True when using CFG

---
 exllamav2/generator/dynamic.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exllamav2/generator/dynamic.py b/exllamav2/generator/dynamic.py
index 14c68858..52924421 100644
--- a/exllamav2/generator/dynamic.py
+++ b/exllamav2/generator/dynamic.py
@@ -1932,7 +1932,7 @@ def emit(
             self.held_k_tokens.append(next_k_tokens)
             self.held_k_probs.append(next_k_probs)
         if self.return_logits:
-            self.held_logits.append(logits)
+            self.held_logits.append(logits[:1, :, :])
 
         # Stop if we reach max_new_tokens