update

examples/rnnoise/step_2_train_model.py

@@ -249,7 +249,6 @@ def main():
 
     step_idx = 0 if last_step_idx == -1 else last_step_idx
 
-    logger.info("training")
     early_stop_flag = False
     for epoch_idx in range(max(0, last_epoch+1), config.max_epochs):
         if early_stop_flag:
@@ -274,6 +273,7 @@ def main():
             noisy_audios: torch.Tensor = noisy_audios.to(device)
 
             denoise_audios, _, _ = model.forward(noisy_audios)
+            denoise_audios = torch.squeeze(denoise_audios, dim=1)
 
             mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
             neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
@@ -339,6 +339,7 @@ def main():
                         noisy_audios: torch.Tensor = noisy_audios.to(device)
 
                         denoise_audios, _, _ = model.forward(noisy_audios)
+                        denoise_audios = torch.squeeze(denoise_audios, dim=1)
 
                         mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
                         neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)