From 2ab9092bc54fd1e0e78a49b9405cf62bb0e57995 Mon Sep 17 00:00:00 2001
From: nlpfollower <nlp@nlpfollower.nlp>
Date: Thu, 23 Jan 2025 13:23:03 -0800
Subject: [PATCH] Decode_n_tokens yield eos/eot token

---
 torchchat/generate.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/torchchat/generate.py b/torchchat/generate.py
index ad933687d..f48543499 100644
--- a/torchchat/generate.py
+++ b/torchchat/generate.py
@@ -576,6 +576,7 @@ def decode_n_tokens(
                         **sampling_kwargs,
                     )
                     input_pos += 1
+                    yield cur_token.clone(), next_prob.clone()
                     break
 
         if not encountered_eos: