| 189 | # TODO: replace with torch layernorm once min required torch version >= 2.1 |
| 190 | class LayerNorm(nn.Module): |
| 191 | def __init__(self, dim, eps: float = 1e-5, elementwise_affine: bool = True, bias: bool = True): |
| 192 | super().__init__() |
| 193 | |
| 194 | self.eps = eps |
| 195 | |
| 196 | if isinstance(dim, numbers.Integral): |
| 197 | dim = (dim,) |
| 198 | |
| 199 | self.dim = torch.Size(dim) |
| 200 | |
| 201 | if elementwise_affine: |
| 202 | self.weight = nn.Parameter(torch.ones(dim)) |
| 203 | self.bias = nn.Parameter(torch.zeros(dim)) if bias else None |
| 204 | else: |
| 205 | self.weight = None |
| 206 | self.bias = None |
| 207 | |
| 208 | def forward(self, input): |
| 209 | return F.layer_norm(input, self.dim, self.weight, self.bias, self.eps) |