Merge branch 'pytorch:main' into newtemplate

pytorch · Feb 8, 2025 · 18ba560 · 18ba560
2 parents 48f8e8b + 53a1004
commit 18ba560
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: