Method
__init__
(
self,
dim: int,
init_values: Union[float, Tensor] = 1e-5,
inplace: bool = False,
device=None,
)
Source from the content-addressed store, hash-verified
| 8 | """Layer scaling module for stable training.""" |
| 9 | |
| 10 | def __init__( |
| 11 | self, |
| 12 | dim: int, |
| 13 | init_values: Union[float, Tensor] = 1e-5, |
| 14 | inplace: bool = False, |
| 15 | device=None, |
| 16 | ) -> None: |
| 17 | super().__init__() |
| 18 | self.inplace = inplace |
| 19 | self.gamma = nn.Parameter(torch.empty(dim, device=device)) |
| 20 | self.init_values = init_values |
| 21 | |
| 22 | def reset_parameters(self): |
| 23 | nn.init.constant_(self.gamma, self.init_values) |
Tested by
no test coverage detected