code: def NMF_separate_spectra(Whats, Xmix):
      """This function separates the mixture signals, given NMF template matrices.
  
      Arguments
      ---------
      Whats : list
          This list contains the list [W1, W2], where W1 W2 are respectively
          the NMF template matrices that correspond to source1 and source2.
          W1, W2 are of size [nfft/2 + 1, K], where nfft is the fft size for STFT,
          and K is the number of vectors (templates) in W.
      Xmix : torch.tensor
          This is the magnitude spectra for the mixtures.
          The size is [BS x T x nfft//2 + 1] where,
          BS = batch size, nfft = fft size, T = number of time steps in the spectra.
  
      Outputs
      -------
      X1hat : Separated spectrum for source1
          Size = [BS x (nfft/2 +1) x T] where,
          BS = batch size, nfft = fft size, T = number of time steps in the spectra.
      X2hat : Separated Spectrum for source2
          The size definitions are the same as above.
  
      Example
      --------
      >>> BS, nfft, T = 4, 20, 400
      >>> K1, K2 = 10, 10
      >>> W1hat = torch.randn(nfft//2 + 1, K1)
      >>> W2hat = torch.randn(nfft//2 + 1, K2)
      >>> Whats = [W1hat, W2hat]
      >>> Xmix = torch.randn(BS, T, nfft//2 + 1)
      >>> X1hat, X2hat = NMF_separate_spectra(Whats, Xmix)
      """
  
      W1, W2 = Whats
  
      nmixtures = Xmix.shape[0]
      Xmix = Xmix.permute(0, 2, 1).reshape(-1, Xmix.size(-1)).t()
      n = Xmix.shape[1]
      eps = 1e-20
  
      # Normalize input
      g = Xmix.sum(dim=0) + eps
      z = Xmix / g
  
      # initialize
      w = torch.cat([W1, W2], dim=1)
      K = w.size(1)
      K1 = W1.size(1)
  
      h = 0.1 * torch.rand(K, n)
      h /= torch.sum(h, dim=0) + eps
  
      for ep in range(1000):
          v = z / (torch.matmul(w, h) + eps)
  
          nh = h * torch.matmul(w.t(), v)
          h = nh / (torch.sum(nh, dim=0) + eps)
  
      h *= g
      Xhat1 = torch.matmul(w[:, :K1], h[:K1, :])
      Xhat1 = torch.split(
          Xhat1.unsqueeze(0), Xhat1.size(1) // nmixtures, dim=2
      )
      Xhat1 = torch.cat(Xhat1, dim=0)
  
      Xhat2 = torch.matmul(w[:, K1:], h[K1:, :])
      Xhat2 = torch.split(
          Xhat2.unsqueeze(0), Xhat2.size(1) // nmixtures, dim=2
      )
      Xhat2 = torch.cat(Xhat2, dim=0)
  
      return Xhat1, Xhat2
  
imports: import torch
  from speechbrain.processing.features import spectral_magnitude
  import speechbrain.processing.features
label_desc: Use image node_id 2 for calling a global function with example usage: NMF_separate_spectra(Whats, Xmix) and returns: Xhat1,  Xhat2