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/conformer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) (limited to 'text_recognizer/networks/conformer/conformer.py') diff --git a/text_recognizer/networks/conformer/conformer.py b/text_recognizer/networks/conformer/conformer.py index 8d0e98e..e2dce27 100644 --- a/text_recognizer/networks/conformer/conformer.py +++ b/text_recognizer/networks/conformer/conformer.py @@ -10,6 +10,8 @@ from text_recognizer.networks.conformer.block import ConformerBlock class Conformer(nn.Module): def __init__( self, + dim: int, + num_classes: int, subsampler: Type[nn.Module], block: ConformerBlock, depth: int, @@ -17,9 +19,10 @@ class Conformer(nn.Module): super().__init__() self.subsampler = subsampler self.blocks = nn.ModuleList([deepcopy(block) for _ in range(depth)]) + self.fc = nn.Linear(dim, num_classes, bias=False) def forward(self, x: Tensor) -> Tensor: x = self.subsampler(x) for fn in self.blocks: x = fn(x) - return x + return self.fc(x) -- cgit v1.2.3-70-g09d2