Add encoded size to start_pos

nlpfollower · nlpfollower · commit 0852578d71fd · 2025-01-15T19:09:53.000-08:00
diff --git a/torchchat/generate.py b/torchchat/generate.py
@@ -1187,6 +1187,7 @@ def callback(x, *, done_generating=False):
                     skip_cache_setup=not is_first_sample,
                     max_seq_length=max_seq_length,
                 )
+                start_pos += encoded.size(0)
                 for token_tensor, metrics in generator_func:
                     if token_tensor is not None:
                         start_pos += token_tensor.size(0)

Original file line number	Diff line number	Diff line change
`@@ -1187,6 +1187,7 @@ def callback(x, *, done_generating=False):`
`1187`	`1187`	`skip_cache_setup=not is_first_sample,`
`1188`	`1188`	`max_seq_length=max_seq_length,`
`1189`	`1189`	`)`
	`1190`	`+ start_pos += encoded.size(0)`
`1190`	`1191`	`for token_tensor, metrics in generator_func:`
`1191`	`1192`	`if token_tensor is not None:`
`1192`	`1193`	`start_pos += token_tensor.size(0)`