update

examples/clean_unet_aishell/step_2_train_model.py

@@ -233,6 +233,9 @@ def main():
             ae_loss = ae_loss_fn(enhanced_audios, clean_audios)
             sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios, clean_audios)
 
+            if torch.any(torch.isnan(mag_loss)) or torch.any(torch.isinf(mag_loss)):
+                raise AssertionError("nan or inf in mag_loss")
+
             loss = ae_loss + sc_loss + mag_loss
 
             enhanced_audios_list_r = list(enhanced_audios.detach().cpu().numpy())