Decode_n_tokens yield eos/eot token

pytorch · Jan 23, 2025 · 2ab9092 · 2ab9092
1 parent 3ce9c8e
commit 2ab9092
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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: