update

examples/clean_unet_aishell/step_2_train_model.py

@@ -227,12 +227,8 @@ def main():
             clean_audios = clean_audios.to(device)
             noisy_audios = noisy_audios.to(device)
 
-            print(f"clean_audios shape: {clean_audios.shape}, dtype: {clean_audios.dtype}")
-            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)
             sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios, clean_audios)
@@ -248,7 +244,7 @@ def main():
             optimizer.step()
             lr_scheduler.step()
 
-            total_pesq_metric += pesq_metric.item()
+            total_pesq_metric += 0 if pesq_metric is None else pesq_metric.item()
             total_loss += loss.item()
             total_ae_loss += ae_loss.item()
             total_sc_loss += sc_loss.item()