Merge pull request espnet#4343 from Emrys365/complex_support

Fix a bug in stats aggregation when PITSolver is used
chintu619 · May 8, 2022 · 537f9b6 · 537f9b6
2 parents 793b999 + 9e8e753
commit 537f9b6
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/espnet2/enh/loss/wrappers/pit_solver.py b/espnet2/enh/loss/wrappers/pit_solver.py
@@ -73,13 +73,14 @@ def pair_loss(permutation):
             )
             # remove stats from unused permutations
             for k, v in stats.items():
-                # (B, len(all_permutations), ...)
+                # (B, num_spk * len(all_permutations), ...)
                 new_v = torch.stack(v, dim=1)
+                B, L, *rest = new_v.shape
+                assert L == num_spk * len(all_permutations), (L, num_spk)
+                new_v = new_v.view(B, L // num_spk, num_spk, *rest).mean(2)
                 if new_v.dim() > 2:
-                    shapes = [1 for _ in range(new_v.dim() - 2)]
-                    perm0 = perm_.view(perm_.shape[0], 1, *shapes).expand(
-                        -1, -1, *new_v.shape[2:]
-                    )
+                    shapes = [1 for _ in rest]
+                    perm0 = perm_.view(perm_.shape[0], 1, *shapes).expand(-1, -1, *rest)
                 else:
                     perm0 = perm_.unsqueeze(1)
                 stats[k] = new_v.gather(1, perm0.to(device=new_v.device)).unbind(1)