| 386 | # TODO: replace with torch layernorm once min required torch version >= 2.1 |
| 387 | class LayerNorm(nn.Module): |
| 388 | def __init__(self, dim, eps: float = 1e-5, elementwise_affine: bool = True, bias: bool = True): |
| 389 | super().__init__() |
| 390 | |
| 391 | self.eps = eps |
| 392 | |
| 393 | if isinstance(dim, numbers.Integral): |
| 394 | dim = (dim,) |
| 395 | |
| 396 | self.dim = torch.Size(dim) |
| 397 | |
| 398 | if elementwise_affine: |
| 399 | self.weight = nn.Parameter(torch.ones(dim)) |
| 400 | self.bias = nn.Parameter(torch.zeros(dim)) if bias else None |
| 401 | else: |
| 402 | self.weight = None |
| 403 | self.bias = None |
| 404 | |
| 405 | def forward(self, input): |
| 406 | return F.layer_norm(input, self.dim, self.weight, self.bias, self.eps) |