From c0a1f78de81a17ab0b441982cca3ca41c67b6411 Mon Sep 17 00:00:00 2001 From: Mohammad Zeineldeen Date: Thu, 4 Nov 2021 12:23:46 +0100 Subject: [PATCH] better use sequential --- nn/conformer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/nn/conformer.py b/nn/conformer.py index 2d0a1fbc..33f542b7 100644 --- a/nn/conformer.py +++ b/nn/conformer.py @@ -187,7 +187,7 @@ def __init__(self, encoder_layer: nn.Module, num_blocks: int, conv_kernel_size: self.linear = nn.Linear(n_out=enc_key_dim, l2=l2, with_bias=False) - self.conformer_blocks = nn.ModuleList([ + self.conformer_blocks = nn.Sequential([ encoder_layer( conv_kernel_size=conv_kernel_size, ff_act=ff_act, ff_dim=ff_dim, dropout=dropout, att_dropout=att_dropout, enc_key_dim=enc_key_dim, att_n_heads=att_n_heads, l2=l2 @@ -199,6 +199,5 @@ def forward(self, inp: LayerRef) -> LayerRef: x_subsample = self.conv_subsample_layer(inp) x_linear = self.linear(x_subsample) x = nn.dropout(x_linear, dropout=self.dropout) - for conformer_block in self.conformer_blocks: - x = conformer_block(x) + x = self.conformer_blocks(x) return x