Fix text length type in TTSDataset for beta_binomial_interpolator (NV…

…IDIA#3233) * fix text length type Signed-off-by: Oktai Tatanov <[email protected]> * fix log mel length Signed-off-by: Oktai Tatanov <[email protected]> * fix log mel length 2 Signed-off-by: Oktai Tatanov <[email protected]> * fix log mel length 3 Signed-off-by: Oktai Tatanov <[email protected]> Co-authored-by: Yang Zhang <[email protected]>
nik-ko · Nov 23, 2021 · 0178ff7 · 0178ff7
1 parent f43480e
commit 0178ff7
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/nemo/collections/tts/torch/data.py b/nemo/collections/tts/torch/data.py
@@ -359,17 +359,17 @@ def __getitem__(self, index):
 
         duration_prior = None
         if DurationPrior in self.sup_data_types_set:
-            log_mel_length = torch.tensor(self.get_log_mel(audio).squeeze(0).shape[1]).long()
-
             if self.use_beta_binomial_interpolator:
-                duration_prior = torch.from_numpy(self.beta_binomial_interpolator(log_mel_length.item(), text_length))
+                mel_len = self.get_log_mel(audio).shape[2]
+                duration_prior = torch.from_numpy(self.beta_binomial_interpolator(mel_len, text_length.item()))
             else:
                 prior_path = Path(self.sup_data_path) / f"pr_{audio_stem}.pt"
 
                 if prior_path.exists():
                     duration_prior = torch.load(prior_path)
                 else:
-                    duration_prior = beta_binomial_prior_distribution(text_length, log_mel_length)
+                    mel_len = self.get_log_mel(audio).shape[2]
+                    duration_prior = beta_binomial_prior_distribution(text_length, mel_len)
                     duration_prior = torch.from_numpy(duration_prior)
                     torch.save(duration_prior, prior_path)