Dynamic gen: Fix return_logits=True when using CFG

turboderp-org · Jul 10, 2024 · 25e46d8 · 25e46d8
1 parent 1179b8a
commit 25e46d8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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