update

examples/spectrum_unet_irm_aishell/run.sh

@@ -12,7 +12,7 @@ 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 \
+sh run.sh --stage 2 --stop_stage 2 --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"
 

examples/spectrum_unet_irm_aishell/step_2_train_model.py

@@ -142,6 +142,9 @@ class CollateFunction(object):
             snr_db: torch.Tensor = 10 * torch.log10(
                 speech_spec / (noise_spec + self.epsilon)
             )
+            if torch.any(torch.isnan(snr_db)) or torch.any(torch.isinf(snr_db)):
+                raise AssertionError("nan or inf in snr_db")
+
             snr_db_ = torch.unsqueeze(snr_db, dim=0)
             snr_db_ = torch.unsqueeze(snr_db_, dim=0)
             snr_db_ = self.make_unfold_snr_db(snr_db_, n_time_steps=3)