From 0bfd81d7213e17090f56912260c3f0054f17ba8c Mon Sep 17 00:00:00 2001
From: Daniel Povey <dpovey@gmail.com>
Date: Mon, 28 Nov 2022 10:42:06 +0800
Subject: [PATCH] fix bug RE dims_to_mean

---
 egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py b/egs/librispeech/ASR/pruned_transducer_stateless7/scaling.py
index 49e204de1..f685bf112 100644
--- 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