Fix (#7478)

Signed-off-by: Cheng-Ping Hsieh <[email protected]>
NVIDIA · Sep 22, 2023 · a9fb58b · a9fb58b
1 parent 433d9f1
commit a9fb58b
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/nemo/collections/nlp/data/language_modeling/megatron/gpt_sft_chat_dataset.py b/nemo/collections/nlp/data/language_modeling/megatron/gpt_sft_chat_dataset.py
@@ -289,6 +289,7 @@ def collate_fn(self, batch):
             labels = [x[: self.max_seq_length] for x in labels]
             loss_mask = [x[: self.max_seq_length] for x in loss_mask]
             contexts = [x[: self.max_seq_length] for x in contexts]
+            answers = [x[: self.max_seq_length] for x in answers]
 
         # increase max length to nearest multiple of 4 or 8
         if self.pad_to_max_length: