| 298 | |
| 299 | |
| 300 | class LayerNorm(nn.LayerNorm): |
| 301 | def __init__(self, *args, **kwargs): |
| 302 | """Initialize LayerNorm. |
| 303 | |
| 304 | Args: |
| 305 | *args: Variable positional arguments. |
| 306 | **kwargs: Additional keyword arguments. |
| 307 | """ |
| 308 | super().__init__(*args, **kwargs) |
| 309 | |
| 310 | def forward(self, input): |
| 311 | """Forward pass for training. |
| 312 | |
| 313 | Args: |
| 314 | input: Input audio/text data. |
| 315 | """ |
| 316 | output = F.layer_norm( |
| 317 | input.float(), |
| 318 | self.normalized_shape, |
| 319 | self.weight.float() if self.weight is not None else None, |
| 320 | self.bias.float() if self.bias is not None else None, |
| 321 | self.eps, |
| 322 | ) |
| 323 | return output.type_as(input) |
| 324 | |
| 325 | |
| 326 | def sequence_mask(lengths, maxlen=None, dtype=torch.float32, device=None): |