update

examples/spectrum_unet_irm_aishell/run.sh

@@ -12,6 +12,10 @@ sh run.sh --stage 1 --stop_stage 3 --system_version centos --file_folder_name fi
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 
+sh run.sh --stage 2 --stop_stage 3 --system_version centos --file_folder_name file_dir \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
+
 
 END
 

examples/spectrum_unet_irm_aishell/step_2_train_model.py

@@ -290,6 +290,10 @@ def main():
             snr_db_target = snr_db.to(device)
 
             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")
             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.01 * snr_loss
@@ -325,6 +329,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")
                 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.01 * snr_loss