update

toolbox/torchaudio/models/spectrum_dfnet/modeling_spectrum_dfnet.py

@@ -849,6 +849,11 @@ class SpectrumDfNet(nn.Module):
         spec_e = torch.squeeze(spec_e, dim=1)
         spec_e = spec_e.permute(0, 2, 1, 3)
         # spec_e shape: [batch_size, spec_bins, time_steps, 2]
+
+        mask = torch.squeeze(mask, dim=1)
+        mask = mask.permute(0, 2, 1)
+        # mask shape: [batch_size, spec_bins, time_steps]
+
         return spec_e, mask, lsnr