update

examples/clean_unet_aishell/step_2_train_model.py

@@ -236,19 +236,9 @@ def main():
             enhanced_audios = model.forward(noisy_audios)
             enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
 
-            if torch.any(torch.isnan(clean_audios)) or torch.any(torch.isinf(clean_audios)):
-                raise AssertionError("nan or inf in clean_audios")
-            if torch.any(torch.isnan(noisy_audios)) or torch.any(torch.isinf(noisy_audios)):
-                raise AssertionError("nan or inf in noisy_audios")
-            if torch.any(torch.isnan(enhanced_audios)) or torch.any(torch.isinf(enhanced_audios)):
-                raise AssertionError("nan or inf in enhanced_audios")
-
             ae_loss = ae_loss_fn(enhanced_audios, clean_audios)
             sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios, clean_audios)
 
-            if torch.any(torch.isnan(mag_loss)) or torch.any(torch.isinf(mag_loss)):
-                raise AssertionError("nan or inf in mag_loss")
-
             loss = ae_loss + sc_loss + mag_loss
 
             enhanced_audios_list_r = list(enhanced_audios.detach().cpu().numpy())

toolbox/torchaudio/models/clean_unet/modeling_clean_unet.py

@@ -144,7 +144,7 @@ class CleanUNet(nn.Module):
                     nn.Conv1d(channels_h, channels_h * 2, 1),
                     nn.GLU(dim=1),
                     nn.ConvTranspose1d(channels_h, channels_output, kernel_size, stride),
-                    nn.ReLU(inplace=False)
+                    # nn.ReLU(inplace=False)
                 ))
             channels_output = channels_h