update

examples/clean_unet_aishell/step_1_prepare_data.py

@@ -101,7 +101,7 @@ def get_dataset(args):
     count = 0
     process_bar = tqdm(desc="build dataset excel")
     for noise, speech in zip(noise_generator, speech_generator):
-        if count > args.max_count:
+        if count >= args.max_count:
             break
 
         noise_filename = noise["filename"]

examples/clean_unet_aishell/step_2_train_model.py

@@ -232,8 +232,8 @@ def main():
 
             loss = ae_loss + sc_loss + mag_loss
 
-            enhanced_audios_list_r = list(enhanced_audios.cpu().numpy())
-            clean_audios_list_r = list(clean_audios.cpu().numpy())
+            enhanced_audios_list_r = list(enhanced_audios.detach().cpu().numpy())
+            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
             pesq_metric = batch_pesq(enhanced_audios_list_r, clean_audios_list_r)
 
             optimizer.zero_grad()
@@ -292,8 +292,8 @@ def main():
 
                 loss = ae_loss + sc_loss + mag_loss
 
-                enhanced_audios_list_r = list(enhanced_audios.cpu().numpy())
-                clean_audios_list_r = list(clean_audios.cpu().numpy())
+                enhanced_audios_list_r = list(enhanced_audios.detach().cpu().numpy())
+                clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
                 pesq_metric = batch_pesq(enhanced_audios_list_r, clean_audios_list_r)
 
                 total_pesq_metric += pesq_metric.item()