From 00c7aeff97eac27f72dfec15c99310b6cc6e9c50 Mon Sep 17 00:00:00 2001 From: "haoxiangsnr@gmail.com" Date: Thu, 28 Dec 2023 11:17:05 +0000 Subject: [PATCH] fix: add enh_mag calculation for DNSMOS loss --- recipes/intel_ndns/spiking_fullsubnet/model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/recipes/intel_ndns/spiking_fullsubnet/model.py b/recipes/intel_ndns/spiking_fullsubnet/model.py index 7b7d442..0790e8e 100644 --- a/recipes/intel_ndns/spiking_fullsubnet/model.py +++ b/recipes/intel_ndns/spiking_fullsubnet/model.py @@ -417,10 +417,11 @@ def forward(self, input): enh_stft = noisy_cmp.clone() enh_stft[..., :-1, :] = enh_freqs enh_stft = rearrange(enh_stft, "b 1 f t -> b f t") + enh_mag = torch.abs(enh_stft) # For computing DNSMOS loss enh_y = self.istft(enh_stft, length=sequence_length) - return enh_y, fb_all_layer_outputs, sb_all_layer_outputs + return enh_y, enh_mag, fb_all_layer_outputs, sb_all_layer_outputs if __name__ == "__main__":