summaryrefslogtreecommitdiff
path: root/text_recognizer/networks/conformer/scale.py
blob: da736a33b9cb9b5b3c94e4b3ad9da4b5da7e10b7 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
"""Scale layer."""
from typing import Dict
from torch import nn, Tensor


class Scale(nn.Module):
    def __init__(self, scale: float, fn: nn.Module) -> None:
        super().__init__()
        self.scale = scale
        self.fn = fn

    def forward(self, x: Tensor, kwargs: Dict) -> Tensor:
        return self.fn(x, **kwargs) * self.scale