Training: No speech after 400K steps

[fakerybakery]

Hi,
I used the default training script for ~400K steps using the attached inference code, but I'm not getting anything that sounds close to speech, does anyone know if I'm doing something wrong with the training or inference?
Thanks!

Inference code

from pathlib import Path

import torch
import torchaudio

from einops import rearrange
from torchaudio.transforms import GriffinLim, InverseSpectrogram, InverseMelScale, Resample, Speed
import matplotlib.pyplot as plt
from e2_tts_pytorch.e2_tts import (
    E2TTS,
    DurationPredictor,
    MelSpec
)
n_fft = 1024
sample_rate = 24_000
from vocos import Vocos

# Move models to CUDA
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

duration_predictor = DurationPredictor(
    transformer=dict(
        dim=512,
        depth=6,
    )
).to(device)

e2tts = E2TTS(
    duration_predictor=duration_predictor,
    transformer=dict(
        dim=512,
        depth=12,
        skip_connect_type='concat'
    )
).to(device)

# Load checkpoint onto GPU
checkpoint = torch.load('./e2tts.pt', map_location=device)
e2tts.load_state_dict(checkpoint['model_state_dict'])

# Text input
text = [
    "It took me quite a long time to develop a voice, and now that I have it I'm not going to be silent.",
    "Waves crashed against the cliffs, their thunderous applause echoing for miles."
]

from torchaudio.transforms import MelSpectrogram
from e2_tts_pytorch.e2_tts import log

audio, sr = torchaudio.load('ref.wav')
audio = audio.to(device)

text = text[0]
print(f"text: {text}")

with torch.inference_mode():
    generated = e2tts.sample(
        cond=audio,
        text=[text],
        steps=32,
        cfg_strength=0.5,
    ).cpu()
    print(f"Generated: {generated.shape}")
    print(generated)
    # If generated tensor has 2 dimensions, reshape it to add the third dimension
    #if generated.dim() == 2:
    #    generated = generated.unsqueeze(-1)  # Add the third dimension

    # Rearrange the tensor to match the expected shape
    #generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')
plt.figure(figsize=(12, 4))
generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')
plt.imshow(generated_mel_spec[0].numpy(), origin='lower', aspect='auto')
plt.colorbar()
plt.savefig('out.png')