1. Debugging.

Signed-off-by: Micha Livne <[email protected]>
NVIDIA · Mar 8, 2022 · 31e536f · 31e536f
1 parent f7223b1
commit 31e536f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/nemo/collections/nlp/data/language_modeling/megatron/bart_dataset.py b/nemo/collections/nlp/data/language_modeling/megatron/bart_dataset.py
@@ -206,7 +206,7 @@ def build_training_sample(
 
     # Masking.
     max_predictions_per_seq = max_num_tokens
-    (tokens, masked_positions, masked_labels, _, masked_spans) = create_masked_lm_predictions(
+    (output_tokens, masked_positions, masked_labels, _, masked_spans) = create_masked_lm_predictions(
         tokens=tokens,
         vocab_id_list=vocab_id_list,
         vocab_id_to_token_dict=vocab_id_to_token_dict,
@@ -229,6 +229,7 @@ def build_training_sample(
     # Padding.
     tokens_enc, tokens_dec_in, labels, enc_mask, dec_mask, loss_mask = pad_and_convert_to_numpy(
         tokens,
+        output_tokens,
         masked_positions,
         masked_labels,
         pad_id,
@@ -255,6 +256,7 @@ def build_training_sample(
 
 def pad_and_convert_to_numpy(
     tokens,
+    output_tokens,
     masked_positions,
     masked_labels,
     pad_id,

diff --git a/nemo/collections/nlp/data/language_modeling/megatron/t5_dataset.py b/nemo/collections/nlp/data/language_modeling/megatron/t5_dataset.py
@@ -266,7 +266,6 @@ def pad_and_convert_to_numpy(
     sentinel_tokens=None,
 ):
     """Pad sequences and convert them to numpy."""
-    import pdb; pdb.set_trace()
     sentinel_tokens = collections.deque(sentinel_tokens)
     t5_input = []
     (t5_decoder_in, t5_decoder_out) = ([bos_id], [])