diff --git a/TTS/tts/layers/losses.py b/TTS/tts/layers/losses.py index c1bc85b5..50809de5 100644 --- a/TTS/tts/layers/losses.py +++ b/TTS/tts/layers/losses.py @@ -402,4 +402,9 @@ class GlowTTSLoss(torch.nn.Module): return_dict['loss'] = log_mle + loss_dur return_dict['log_mle'] = log_mle return_dict['loss_dur'] = loss_dur + + # check if any loss is NaN + for key, loss in return_dict.items(): + if torch.isnan(loss): + raise RuntimeError(f" [!] NaN loss with {key}.") return return_dict \ No newline at end of file