fix: gradient truncation bug in pit_solver.py

Masao-Someki · Dec 7, 2024 · d5382f1 · d5382f1
1 parent cccc290
commit d5382f1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/espnet2/enh/loss/wrappers/pit_solver.py b/espnet2/enh/loss/wrappers/pit_solver.py
@@ -98,9 +98,9 @@ def pair_loss(permutation):
                     perm0 = perm_.unsqueeze(1)
                 stats[k] = new_v.gather(1, perm0.to(device=new_v.device)).unbind(1)
         else:
-            loss = torch.tensor(
+            loss = torch.stack(
                 [
-                    torch.tensor(
+                    torch.stack(
                         [
                             pre_hook(
                                 self.criterion,