text_recognizer/networks/conformer/conformer.py


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

"""Conformer module."""
from copy import deepcopy

from torch import nn, Tensor

from text_recognizer.networks.conformer.block import ConformerBlock


class Conformer(nn.Module):
    def __init__(self, block: ConformerBlock, depth: int) -> None:
        super().__init__()
        self.blocks = nn.ModuleList([deepcopy(block) for _ in range(depth)])

    def forward(self, x: Tensor) -> Tensor:
        for fn in self.blocks:
            x = fn(x)
        return x