fix bug RE dims_to_mean

2022-11-28 10:42:06 +08:00 · 2022-11-28 10:42:06 +08:00 · 0bfd81d721
commit 0bfd81d721
parent 109825cafb
1 changed files with 1 additions and 0 deletions
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py
@ -422,6 +422,7 @@ class LinearWithAuxLossFunction(torch.autograd.Function):
                y = torch.matmul(x, weight.t())
                z = torch.matmul(y, weight)
                # subtract mean
+                dims_to_mean = tuple(range(x.ndim-1))
                x = x - x.mean(dim=dims_to_mean)
                z = z - z.mean(dim=dims_to_mean)
                # compute optimal scale on z