[TTS] enhancer dataset: correct dim for length (NVIDIA#6048) (NVIDIA#…

…6074) Signed-off-by: Roman Korostik <[email protected]> Co-authored-by: Roman Korostik <[email protected]>
titu1994 · Mar 24, 2023 · 0e6322f · 0e6322f
1 parent c583f13
commit 0e6322f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/tts/torch/data.py b/nemo/collections/tts/torch/data.py
@@ -1136,7 +1136,7 @@ def __getitem__(self, index):
 
     def _collate_fn(self, batch):
         pred_specs, true_specs = zip(*batch)
-        lengths = [spec.shape[-1] for spec in true_specs]
+        lengths = [spec.shape[0] for spec in true_specs]
 
         pred_specs = torch.nn.utils.rnn.pad_sequence(pred_specs, batch_first=True)
         true_specs = torch.nn.utils.rnn.pad_sequence(true_specs, batch_first=True)