update

examples/clean_unet_aishell/step_2_train_model.py

@@ -229,6 +229,8 @@ def main():
 
             enhanced_audios = model.forward(noisy_audios)
             enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
+            if torch.any(torch.isnan(enhanced_audios)) or torch.any(torch.isinf(enhanced_audios)):
+                raise AssertionError("nan or inf in enhanced_audios")
 
             ae_loss = ae_loss_fn(enhanced_audios, clean_audios)
             sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios, clean_audios)