bug fix in tarred dataset class num_speech_codebooks

Signed-off-by: Paarth Neekhara <[email protected]>
paarthneekhara · Dec 14, 2023 · 7d985ca · 7d985ca
1 parent 9c3d712
commit 7d985ca
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
@@ -943,7 +943,7 @@ def _build_sample(self, tup):
                     dim=1,
                 )
                 dec_input_new = []
-                for _c in range(self.num_codebooks):
+                for _c in range(self.num_speech_codebooks):
                     st = num_codebooks - _c
                     et_decoder_input = dec_input_padded.shape[1] - _c
                     dec_input_new.append(dec_input_padded[_c, st:et_decoder_input])