Merge pull request #12 from wetdog/main

Fix argument in duration predictor
lucidrains · Jul 22, 2024 · 15c4de2 · 15c4de2
2 parents d8faa5e + 73f5c87
commit 15c4de2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/e2_tts_pytorch/trainer.py b/e2_tts_pytorch/trainer.py
@@ -203,7 +203,7 @@ def train(self, train_dataset, epochs, batch_size, grad_accumulation_steps=1, nu
                 mel_lengths = batch["mel_lengths"]
 
                 if self.duration_predictor is not None:
-                    dur_loss = self.duration_predictor(mel_spec, target_duration=batch.get('durations'))
+                    dur_loss = self.duration_predictor(mel_spec, lens=batch.get('durations'))
                     self.writer.add_scalar('duration loss', dur_loss.item(), global_step)
 
                 loss = self.model(mel_spec, text=text_inputs, lens=mel_lengths)