| 266 | |
| 267 | |
| 268 | class LayerNorm(nn.LayerNorm): |
| 269 | def __init__(self, *args, **kwargs): |
| 270 | super().__init__(*args, **kwargs) |
| 271 | |
| 272 | def forward(self, input): |
| 273 | output = F.layer_norm( |
| 274 | input.float(), |
| 275 | self.normalized_shape, |
| 276 | self.weight.float() if self.weight is not None else None, |
| 277 | self.bias.float() if self.bias is not None else None, |
| 278 | self.eps, |
| 279 | ) |
| 280 | return output.type_as(input) |
| 281 | |
| 282 | |
| 283 | def sequence_mask(lengths, maxlen=None, dtype=torch.float32, device=None): |