update

examples/conv_tasnet/step_2_train_model.py

@@ -141,7 +141,7 @@ def main():
         max_wave_value=32768.0,
         min_snr_db=config.min_snr_db,
         max_snr_db=config.max_snr_db,
-        skip=675000,
+        skip=825000,
     )
     valid_dataset = DenoiseJsonlDataset(
         jsonl_file=args.valid_dataset,
@@ -296,6 +296,9 @@ def main():
             # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss
             # loss = 2.0 * mr_stft_loss + 0.8 * ae_loss + 0.7 * neg_si_snr_loss + 0.5 * neg_stoi_loss
             loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss + 0.5 * pesq_loss
+            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                logger.info(f"find nan or inf in loss.")
+                continue
 
             denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
             clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
@@ -374,6 +377,9 @@ def main():
                         # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss
                         # loss = 2.0 * mr_stft_loss + 0.8 * ae_loss + 0.7 * neg_si_snr_loss + 0.5 * neg_stoi_loss
                         loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss + 0.5 * pesq_loss
+                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                            logger.info(f"find nan or inf in loss.")
+                            continue
 
                         denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
                         clean_audios_list_r = list(clean_audios.detach().cpu().numpy())

requirements.txt

@@ -1,4 +1,4 @@
-gradio
+gradio==5.23.2
 datasets==3.2.0
 python-dotenv==1.0.1
 scipy==1.15.1
@@ -12,4 +12,4 @@ torch-pesq==0.1.2
 torchmetrics==1.6.1
 torchmetrics[audio]==1.6.1
 einops==0.8.1
-torch_stoi==0.2.3
+torch-stoi==0.2.3