Merge branch 'multiple-eos-edge-case' into 'master'

Fixed edge case with multiple end of sequence in one sequence See merge request ADLR/megatron-lm!12

Merge branch 'multiple-eos-edge-case' into 'master'
Fixed edge case with multiple end of sequence in one sequence See merge request ADLR/megatron-lm!12
0f873f97 · Mohammad Shoeybi · cf2f4d9d · 872b4a6a · 0f873f97
Commit 0f873f97 authored Nov 04, 2019 by Mohammad Shoeybi
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