fix

train_ms_mel.py

@@ -134,21 +134,47 @@ def unpack_batch(
     hps: HyperParameters,
     device: torch.device,
 ) -> tuple[dict[str, Any], torch.Tensor]:
-    (
+    batch = tuple(item.to(device, non_blocking=True) for item in batch)
-        x,
+    if hps.data.use_jp_extra:
-        x_lengths,
+        (
-        spec,
+            x,
-        spec_lengths,
+            x_lengths,
-        waveform,
+            spec,
-        _,
+            spec_lengths,
-        speakers,
+            waveform,
-        tone,
+            _,
-        language,
+            speakers,
-        bert,
+            tone,
-        ja_bert,
+            language,
-        en_bert,
+            bert,
-        style_vec,
+            style_vec,
-    ) = (item.to(device, non_blocking=True) for item in batch)
+        ) = batch
+        encoder_args = (tone, language, bert, style_vec)
+    else:
+        (
+            x,
+            x_lengths,
+            spec,
+            spec_lengths,
+            waveform,
+            _,
+            speakers,
+            tone,
+            language,
+            bert,
+            ja_bert,
+            en_bert,
+            style_vec,
+        ) = batch
+        encoder_args = (
+            tone,
+            language,
+            bert,
+            ja_bert,
+            en_bert,
+            style_vec,
+            speakers,
+        )
     mel = (
         spec
         if hps.model.use_mel_posterior_encoder
@@ -161,11 +187,6 @@ def unpack_batch(
             hps.data.mel_fmax,
         )
     )
-    encoder_args = (
-        (tone, language, bert, style_vec)
-        if hps.data.use_jp_extra
-        else (tone, language, bert, ja_bert, en_bert, style_vec, speakers)
-    )
     return {
         "x": x,
         "x_lengths": x_lengths,
@@ -260,7 +281,7 @@ def run() -> None:
         shuffle=True,
         num_workers=args.num_workers,
         pin_memory=device.type == "cuda",
-        collate_fn=TextAudioSpeakerCollate(),
+        collate_fn=TextAudioSpeakerCollate(use_jp_extra=hps.data.use_jp_extra),
         drop_last=True,
     )
     output_dir = Path(args.output_dir or Path(args.config).parent / "models")