update

examples/spectrum_unet_irm_aishell/step_2_train_model.py

@@ -168,11 +168,11 @@ class CollateFunction(object):
 
         # assert
         if torch.any(torch.isnan(mix_spec_list)) or torch.any(torch.isinf(mix_spec_list)):
-            raise AssertionError("nan in mix_spec_list")
+            raise AssertionError("nan or inf in mix_spec_list")
         if torch.any(torch.isnan(speech_irm_list)) or torch.any(torch.isinf(speech_irm_list)):
-            raise AssertionError("nan in speech_irm_list")
+            raise AssertionError("nan or inf in speech_irm_list")
         if torch.any(torch.isnan(snr_db_list)) or torch.any(torch.isinf(snr_db_list)):
-            raise AssertionError("nan in snr_db_list")
+            raise AssertionError("nan or inf in snr_db_list")
 
         return mix_spec_list, speech_irm_list, snr_db_list
 
@@ -291,7 +291,7 @@ def main():
 
             speech_irm_prediction, lsnr_prediction = model.forward(mix_spec)
             if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
-                raise AssertionError("nan in speech_irm_prediction")
+                raise AssertionError("nan or inf in speech_irm_prediction")
             if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
                 raise AssertionError("nan or inf in lsnr_prediction")
             irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
@@ -331,10 +331,10 @@ def main():
 
             with torch.no_grad():
                 speech_irm_prediction, lsnr_prediction = model.forward(mix_spec)
-                if torch.any(torch.isnan(speech_irm_prediction)):
-                    raise AssertionError("nan in speech_irm_prediction")
-                if torch.any(torch.isnan(lsnr_prediction)):
-                    raise AssertionError("nan in lsnr_prediction")
+                if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
+                    raise AssertionError("nan or inf in speech_irm_prediction")
+                if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
+                    raise AssertionError("nan or inf in lsnr_prediction")
                 irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
                 # snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
                 # loss = irm_loss + 0*snr_loss