Add subsampler layer

1 files changed, 9 insertions, 1 deletions

diff --git a/text_recognizer/networks/conformer/conformer.py b/text_recognizer/networks/conformer/conformer.py
index d56955e..8d0e98e 100644
--- a/text_recognizer/networks/conformer/conformer.py
+++ b/text_recognizer/networks/conformer/conformer.py
@@ -1,5 +1,6 @@
 """Conformer module."""
 from copy import deepcopy
+from typing import Type
 
 from torch import nn, Tensor
 
@@ -7,11 +8,18 @@ from text_recognizer.networks.conformer.block import ConformerBlock
 
 
 class Conformer(nn.Module):
-    def __init__(self, block: ConformerBlock, depth: int) -> None:
+    def __init__(
+        self,
+        subsampler: Type[nn.Module],
+        block: ConformerBlock,
+        depth: int,
+    ) -> None:
         super().__init__()
+        self.subsampler = subsampler
         self.blocks = nn.ModuleList([deepcopy(block) for _ in range(depth)])
 
     def forward(self, x: Tensor) -> Tensor:
+        x = self.subsampler(x)
         for fn in self.blocks:
             x = fn(x)
         return x