update

examples/spectrum_unet_irm_aishell/step_2_train_model.py

@@ -296,6 +296,8 @@ def main():
                 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)
+            if torch.any(torch.isnan(snr_loss)):
+                raise AssertionError("nan in snr_loss")
             loss = irm_loss + 0.00001 * snr_loss
             # loss = irm_loss
 

toolbox/torchaudio/models/spectrum_unet_irm/modeling_spectrum_unet_irm.py

@@ -396,7 +396,6 @@ class Encoder(nn.Module):
         lsnr_ = self.lsnr_fc(emb)
         print(f"lsnr_: {torch.any(torch.isnan(lsnr_))}")
         lsnr = lsnr_ * self.lsnr_scale + self.lsnr_offset
-        print(f"lsnr: {torch.any(torch.isnan(lsnr))}")
         return e0, e1, e2, e3, emb, lsnr