update

toolbox/torchaudio/models/nx_clean_unet/modeling_nx_clean_unet.py

@@ -181,6 +181,8 @@ class NXCleanUNet(nn.Module):
         )
 
     def forward(self, noisy_audios: torch.Tensor):
+        # noisy_audios shape: [batch_size, n_samples]
+        noisy_audios = torch.unsqueeze(noisy_audios, dim=1)
         # noisy_audios shape: [batch_size, 1, n_samples]
 
         n_samples = noisy_audios.shape[-1]
@@ -208,6 +210,9 @@ class NXCleanUNet(nn.Module):
 
         enhanced_audios = enhanced_audios[:, :, :n_samples]
         # enhanced_audios shape: [batch_size, 1, n_samples]
+
+        enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
+
         return enhanced_audios
 
 
@@ -310,11 +315,11 @@ def main():
 
     # shape: [batch_size, channels, num_samples]
     # min length: 94, stride: 32, 32 == 2**5
-    # x = torch.ones([4, 1, 94])
-    # x = torch.ones([4, 1, 126])
-    # x = torch.ones([4, 1, 158])
-    # x = torch.ones([4, 1, 190])
-    x = torch.ones([4, 1, 16000])
+    # x = torch.ones([4, 94])
+    # x = torch.ones([4, 126])
+    # x = torch.ones([4, 158])
+    # x = torch.ones([4, 190])
+    x = torch.ones([4, 16000])
 
     model = NXCleanUNet(config)
     enhanced_audios = model.forward(x)