Fixed edge case with multiple end of sequence in one sequence

872b4a6a · Raul Puri · cf2f4d9d · 872b4a6a
Commit 872b4a6a authored Nov 04, 2019 by Raul Puri
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

generate_samples.py generate_samples.py +1 -1

No files found.
--- a/generate_samples.py
+++ b/generate_samples.py
@@ -417,7 +417,7 @@ def sample_sequence_batch(model, context_tokens, context_lengths, attention_mask
            context_length += 1
            counter += 1

-            done_token = (prev == eos_id).byte()
+            done_token = (prev == eos_id).byte() & started.byte()
            just_finished = (done_token & ~is_done).bool()
            lengths[just_finished.view(-1)] = context_length
            was_done = is_done