Update egs/visinger2/models.py

egs/visinger2/models.py

@@ -573,22 +573,28 @@ class Generator_Noise(torch.nn.Module):
     def forward(self, x, mask):
         istft_x = x
         istft_x = self.istft_pre(istft_x)
-
+    
+        # Apply mask
         istft_x = self.net(istft_x) * mask
-
+    
+        # Compute amplitude and random phase
         amp = self.istft_amplitude(istft_x).unsqueeze(-1)
-        phase = (torch.rand(amp.shape) * 2 * 3.14 - 3.14).to(amp)
-
+        phase = (torch.rand(amp.shape) * 2 * 3.14 - 3.14).to(amp.device)
+    
+        # Calculate real and imaginary parts
         real = amp * torch.cos(phase)
         imag = amp * torch.sin(phase)
-        #spec = torch.cat([real, imag], 3)
+    
+        # Create a complex tensor from real and imaginary parts
         spec = torch.complex(real, imag)
-
-        istft_x = torch.istft(spec, self.fft_size, self.hop_size, self.win_size, self.window.to(amp), True, length=x.shape[2] * self.hop_size, return_complex=False)
-
+    
+        # Inverse short-time Fourier transform
+        istft_x = torch.istft(spec, self.fft_size, self.hop_size, self.win_size, self.window.to(spec.device), True, length=x.shape[2] * self.hop_size, return_complex=False)
+    
         return istft_x.unsqueeze(1)
 
 
+
 class LayerNorm(nn.Module):
     def __init__(self, channels, eps=1e-5):
         super().__init__()