Initialize the RMSNorm normalization layer. Args: dim (int): The dimension of the input tensor. eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. Attributes: eps (float): A small valu
(self, dim: int, eps: float = 1e-6)
| 32 | |
| 33 | class RMSNorm(torch.nn.Module): |
| 34 | def __init__(self, dim: int, eps: float = 1e-6): |
| 35 | """ |
| 36 | Initialize the RMSNorm normalization layer. |
| 37 | |
| 38 | Args: |
| 39 | dim (int): The dimension of the input tensor. |
| 40 | eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. |
| 41 | |
| 42 | Attributes: |
| 43 | eps (float): A small value added to the denominator for numerical stability. |
| 44 | weight (nn.Parameter): Learnable scaling parameter. |
| 45 | |
| 46 | """ |
| 47 | super().__init__() |
| 48 | self.eps = eps |
| 49 | self.weight = nn.Parameter(torch.ones(dim)) |
| 50 | |
| 51 | def _norm(self, x): |
| 52 | """ |