From 7b660c13ce3c0edeace1107838e62c559bc6f078 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Wed, 8 Jun 2022 08:41:09 +0200 Subject: Fix conformer net --- text_recognizer/networks/conformer/subsampler.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) (limited to 'text_recognizer/networks/conformer/subsampler.py') diff --git a/text_recognizer/networks/conformer/subsampler.py b/text_recognizer/networks/conformer/subsampler.py index 2bc0445..53928f1 100644 --- a/text_recognizer/networks/conformer/subsampler.py +++ b/text_recognizer/networks/conformer/subsampler.py @@ -34,13 +34,11 @@ class Subsampler(nn.Module): ) ) subsampler.append(nn.Mish(inplace=True)) - projector = nn.Sequential( - nn.Flatten(start_dim=2), nn.Linear(channels, channels), nn.Dropout(dropout) - ) + projector = nn.Sequential(nn.Linear(channels, channels), nn.Dropout(dropout)) return nn.Sequential(*subsampler), projector def forward(self, x: Tensor) -> Tensor: x = self.subsampler(x) x = self.pixel_pos_embedding(x) - x = self.projector(x) - return x.permute(0, 2, 1) + x = x.flatten(start_dim=2).permute(0, 2, 1) + return self.projector(x) -- cgit v1.2.3-70-g09d2