Invoke the out_balancer of attention_squeeze

2022-12-12 23:31:22 +08:00 · 2022-12-12 23:31:22 +08:00 · b5e0676f14
commit b5e0676f14
parent 0522425ea8
1 changed files with 1 additions and 0 deletions
--- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
+++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py
@ -1445,6 +1445,7 @@ attn_weights: a Tensor of shape (num_heads, batch_size, seq_len, seq_len)
        x = x * scales
        x = self.activation(x)  # Identity only.  For diagnostics.
        x = self.out_proj(x)
+        x = self.out_balancer(x)
        return x