Fix conformer net

1 files changed, 4 insertions, 1 deletions

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)