update

examples/clean_unet_aishell/step_2_train_model.py

@@ -229,6 +229,7 @@ def main():
             print(f"noisy_audios shape: {noisy_audios.shape}, dtype: {noisy_audios.dtype}")
 
             enhanced_audios = model.forward(noisy_audios)
+            enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
             print(f"enhanced_audios shape: {enhanced_audios.shape}, dtype: {enhanced_audios.dtype}")
 
             ae_loss = ae_loss_fn(enhanced_audios, clean_audios)
@@ -291,6 +292,7 @@ def main():
 
                 enhanced_audios = model.forward(noisy_audios)
                 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))