From 31b2a735b8344abe6f73d952fed60f30b3115686 Mon Sep 17 00:00:00 2001 From: Daniel Povey Date: Sun, 20 Nov 2022 13:17:39 +0800 Subject: [PATCH] Move feedforward1 to the beginning, separating it from small_conv_module. --- egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py index f22a8a39a..ff9468167 100644 --- a/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py +++ b/egs/librispeech/ASR/pruned_transducer_stateless7/zipformer.py @@ -472,6 +472,8 @@ class ZipformerEncoderLayer(nn.Module): # multi-headed self-attention module use_self_attn = (random.random() >= dynamic_skip_rate) + src = src + self.feed_forward1(src) + if torch.jit.is_scripting() or use_self_attn: # attn_weights: (num_heads, batch_size, seq_len, seq_len) attn_weights = self.self_attn_weights( @@ -489,8 +491,6 @@ class ZipformerEncoderLayer(nn.Module): if torch.jit.is_scripting() or random.random() >= dynamic_skip_rate: src = src + self.small_conv_module(src, src_key_padding_mask=src_key_padding_mask) - src = src + self.feed_forward1(src) - # pooling module if torch.jit.is_scripting() or use_self_attn: src = src + self.attention_squeeze1(src, attn_weights[1:2])