-1 for not masking first token loss

Signed-off-by: Paarth Neekhara <[email protected]>
paarthneekhara · Nov 23, 2023 · c45e497 · c45e497
1 parent 7e24168
commit c45e497
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/data/language_modeling/megatron/t5_speechlm_tarred_dataset.py b/nemo/collections/nlp/data/language_modeling/megatron/t5_speechlm_tarred_dataset.py
@@ -977,7 +977,7 @@ def collate_fn(self, batch):
             decoder_input_list.append(decoder_input)
             decoder_labels_list.append(decoder_labels)
 
-            decoder_mask[i, :context_tokens_len+question_tokens_len] = 0  # Mask out context and question
+            decoder_mask[i, :context_tokens_len+question_tokens_len-1] = 0  # Mask out context and question
             speech_input_mask[i, context_tokens_len:context_tokens_len+question_tokens_len] = 0
 
         # Using causal attention mask for whole input