[bugfix] EPS summed just before the sqrt operation (#14)

mpariente · Sep 25, 2024 · 2943ec2 · 2943ec2
1 parent 16a7f52
commit 2943ec2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_stoi/stoi.py b/torch_stoi/stoi.py
@@ -156,8 +156,8 @@ def forward(self, est_targets: torch.Tensor,
         #         mask = mask[~not_enough]
 
         # Apply OB matrix to the spectrograms as in Eq. (1)
-        x_tob = torch.matmul(self.OBM, x_spec.abs().pow(2) + EPS).sqrt()
-        y_tob = torch.matmul(self.OBM, y_spec.abs().pow(2) + EPS).sqrt()
+        x_tob = (torch.matmul(self.OBM, x_spec.abs().pow(2)) + EPS).sqrt()
+        y_tob = (torch.matmul(self.OBM, y_spec.abs().pow(2)) + EPS).sqrt()
 
         # Perform N-frame segmentation --> (batch, 15, N, n_chunks)
         x_seg = unfold(x_tob.unsqueeze(2),