fix: forgot to multiply magnitude

test_xumx_2.py

@@ -88,7 +88,8 @@ def separate(
     X = stft(audio_torch)
 
     with torch.no_grad():
-        masked_tf_rep = x_umx_target(X.abs().unsqueeze(0)).squeeze()
+        X_mag = X.abs()
+        masked_tf_rep = x_umx_target(X_mag.unsqueeze(0)).squeeze() * X_mag
 
     Vj = masked_tf_rep.cpu().numpy()
     if softmask: