Spaces:

qgyd2021
/

nx_denoise

Running

HoneyTian commited on Feb 25

Commit

1042eee

1 Parent(s): 46cf9fb

update

Files changed (1) hide show

examples/clean_unet_aishell/step_2_train_model.py CHANGED Viewed

@@ -233,7 +233,7 @@ def main():
             print(f"enhanced_audios shape: {enhanced_audios.shape}, dtype: {enhanced_audios.dtype}")
             ae_loss = ae_loss_fn(enhanced_audios, clean_audios)
-            sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios.squeeze(1), clean_audios.squeeze(1))
             loss = ae_loss + sc_loss + mag_loss
@@ -294,7 +294,7 @@ def main():
                 enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
                 ae_loss = ae_loss_fn(enhanced_audios, enhanced_audios)
-                sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios.squeeze(1), clean_audios.squeeze(1))
                 loss = ae_loss + sc_loss + mag_loss

             print(f"enhanced_audios shape: {enhanced_audios.shape}, dtype: {enhanced_audios.dtype}")
             ae_loss = ae_loss_fn(enhanced_audios, clean_audios)
+            sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios, clean_audios)
             loss = ae_loss + sc_loss + mag_loss
                 enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
                 ae_loss = ae_loss_fn(enhanced_audios, enhanced_audios)
+                sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios, clean_audios)
                 loss = ae_loss + sc_loss + mag_loss