update

examples/spectrum_unet_irm_aishell/step_2_train_model.py

@@ -295,9 +295,9 @@ def main():
             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
-            loss = irm_loss
+            snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
+            loss = irm_loss + 0.00001 * snr_loss
+            # loss = irm_loss
 
             total_loss += loss.item()
             total_examples += mix_spec.size(0)
@@ -334,9 +334,9 @@ def main():
                 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
-                loss = irm_loss
+                snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
+                loss = irm_loss + 0.00001 * snr_loss
+                # loss = irm_loss
 
             total_loss += loss.item()
             total_examples += mix_spec.size(0)