update

examples/spectrum_unet_irm_aishell/step_2_train_model.py

@@ -138,6 +138,10 @@ class CollateFunction(object):
 
             # noise_spec, speech_spec, mix_spec, speech_irm
             # shape: [freq_dim, time_steps]
+            if torch.any(torch.isnan(speech_spec)) or torch.any(torch.isinf(speech_spec)):
+                raise AssertionError("nan or inf in speech_spec")
+            if torch.any(torch.isnan(noise_spec)) or torch.any(torch.isinf(noise_spec)):
+                raise AssertionError("nan or inf in noise_spec")
 
             snr_db: torch.Tensor = 10 * torch.log10(
                 speech_spec / (noise_spec + self.epsilon)