update

toolbox/torchaudio/models/frcrn/inference_frcrn.py

@@ -61,7 +61,9 @@ class InferenceFRCRN(object):
 
         # noisy_audio shape: [batch_size, n_samples]
         enhanced_audio = self.enhancement_by_tensor(noisy_audio)
-        # noisy_audio shape: [n_samples,]
+        # enhanced_audio shape: [channels, num_samples]
+        enhanced_audio = enhanced_audio[0]
+        # enhanced_audio shape: [num_samples]
         return enhanced_audio.cpu().numpy()
 
     def enhancement_by_tensor(self, noisy_audio: torch.Tensor) -> torch.Tensor:
@@ -79,8 +81,7 @@ class InferenceFRCRN(object):
         # shape: [batch_size, 1, num_samples]
 
         enhanced_audio = enhanced_audio[0]
-
-        # enhanced_audio shape: [channels, num_samples]
+        # shape: [channels, num_samples]
         return enhanced_audio
 
 

toolbox/torchaudio/models/mpnet/inference_mpnet.py

@@ -60,7 +60,9 @@ class InferenceMPNet(object):
 
         # noisy_audio shape: [batch_size, n_samples]
         enhanced_audio = self.enhancement_by_tensor(noisy_audio)
-        # noisy_audio shape: [n_samples,]
+        # enhanced_audio shape: [channels, num_samples]
+        enhanced_audio = enhanced_audio[0]
+        # enhanced_audio shape: [num_samples]
         return enhanced_audio.cpu().numpy()
 
     def enhancement_by_tensor(self, noisy_audio: torch.Tensor) -> torch.Tensor:
@@ -81,7 +83,12 @@ class InferenceMPNet(object):
             )
             enhanced_audio = audio_g.detach()
 
+        # shape: [batch_size, num_samples]
+        enhanced_audio = torch.unsqueeze(enhanced_audio, dim=1)
+        # shape: [batch_size, 1, num_samples]
+
         enhanced_audio = enhanced_audio[0]
+        # shape: [channels, num_samples]
         return enhanced_audio