update

toolbox/torchaudio/models/mpnet/inference_mpnet.py

@@ -59,10 +59,9 @@ class InferenceMPNet(object):
         noisy_audio = noisy_audio.unsqueeze(dim=0)
 
         # noisy_audio shape: [batch_size, n_samples]
-        noisy_audio = self.enhancement_by_tensor(noisy_audio)
-        noisy_audio = noisy_audio[0]
-
-        return noisy_audio.cpu().numpy()
+        enhanced_audio = self.enhancement_by_tensor(noisy_audio)
+        # noisy_audio shape: [n_samples,]
+        return enhanced_audio.cpu().numpy()
 
     def enhancement_by_tensor(self, noisy_audio: torch.Tensor) -> torch.Tensor:
         if torch.max(noisy_audio) > 1 or torch.min(noisy_audio) < -1:
@@ -82,6 +81,7 @@ class InferenceMPNet(object):
             )
             enhanced_audio = audio_g.detach()
 
+        enhanced_audio = enhanced_audio[0]
         return enhanced_audio
 
 def main():