update

toolbox/torchaudio/models/spectrum_unet_irm/modeling_spectrum_unet_irm.py

@@ -392,10 +392,9 @@ class Encoder(nn.Module):
         emb = emb.flatten(2)
         # emb shape: [batch_size, time_steps, hidden_size * channels]
         emb, h = self.emb_gru.forward(emb, hidden_state)
-
-        lsnr_ = self.lsnr_fc(emb)
-        print(f"lsnr_: {torch.any(torch.isnan(lsnr_))}")
-        lsnr = lsnr_ * self.lsnr_scale + self.lsnr_offset
+        print(f"emb: {torch.any(torch.isnan(emb))}")
+        lsnr = self.lsnr_fc(emb) * self.lsnr_scale + self.lsnr_offset
+        print(f"lsnr: {torch.any(torch.isnan(lsnr))}")
         return e0, e1, e2, e3, emb, lsnr