update

toolbox/torchaudio/models/mpnet/discriminator.py

@@ -45,7 +45,6 @@ def metric_loss(metric_ref, metrics_gen):
 class MetricDiscriminator(nn.Module):
     def __init__(self, config: MPNetConfig):
         super(MetricDiscriminator, self).__init__()
-
         dim = config.discriminator_dim
         in_channel = config.discriminator_in_channel
 
@@ -86,6 +85,7 @@ class MetricDiscriminatorPretrainedModel(MetricDiscriminator):
         super(MetricDiscriminatorPretrainedModel, self).__init__(
             config=config,
         )
+        self.config = config
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):

toolbox/torchaudio/models/mpnet/modeling_mpnet.py

@@ -151,7 +151,6 @@ class TSTransformerBlock(nn.Module):
 class MPNet(nn.Module):
     def __init__(self, config: MPNetConfig, num_tsblocks=4):
         super(MPNet, self).__init__()
-        self.config = config
         self.num_tscblocks = num_tsblocks
         self.dense_encoder = DenseEncoder(config, in_channel=2)
 
@@ -193,6 +192,7 @@ class MPNetPretrainedModel(MPNet):
         super(MPNetPretrainedModel, self).__init__(
             config=config,
         )
+        self.config = config
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
@@ -251,7 +251,7 @@ def pesq_score(utts_r, utts_g, h):
     pesq_score = Parallel(n_jobs=30)(delayed(eval_pesq)(
                             utts_r[i].squeeze().cpu().numpy(),
                             utts_g[i].squeeze().cpu().numpy(),
-                            h.sample_rate)
+                            h.sample_rate, )
                           for i in range(len(utts_r)))
     pesq_score = np.mean(pesq_score)
 
@@ -260,7 +260,8 @@ def pesq_score(utts_r, utts_g, h):
 
 def eval_pesq(clean_utt, esti_utt, sr):
     try:
-        pesq_score = pesq(sr, clean_utt, esti_utt)
+        mode = "nb" if sr == 8000 else "wb"
+        pesq_score = pesq(sr, clean_utt, esti_utt, mode=mode)
     except:
         pesq_score = -1