(self, dim, eps: float, elementwise_affine: bool = True)
| 408 | |
| 409 | class RMSNorm(nn.Module): |
| 410 | def __init__(self, dim, eps: float, elementwise_affine: bool = True): |
| 411 | super().__init__() |
| 412 | |
| 413 | self.eps = eps |
| 414 | |
| 415 | if isinstance(dim, numbers.Integral): |
| 416 | dim = (dim,) |
| 417 | |
| 418 | self.dim = torch.Size(dim) |
| 419 | |
| 420 | if elementwise_affine: |
| 421 | self.weight = nn.Parameter(torch.ones(dim)) |
| 422 | else: |
| 423 | self.weight = None |
| 424 | |
| 425 | def forward(self, hidden_states): |
| 426 | input_dtype = hidden_states.dtype |