Fix a bug in stats aggregation when PITSolver is used

chintu619 · May 7, 2022 · 5ea4e08 · 5ea4e08
1 parent 793b999
commit 5ea4e08
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/espnet2/enh/loss/wrappers/pit_solver.py b/espnet2/enh/loss/wrappers/pit_solver.py
@@ -73,12 +73,15 @@ 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)]
+                    shapes = [1 for _ in rest]
                     perm0 = perm_.view(perm_.shape[0], 1, *shapes).expand(
-                        -1, -1, *new_v.shape[2:]
+                        -1, -1, *rest
                     )
                 else:
                     perm0 = perm_.unsqueeze(1)