| 22 | |
| 23 | |
| 24 | class LayerScale(nn.Module): |
| 25 | |
| 26 | def __init__(self, |
| 27 | dim: int, |
| 28 | inplace: bool = False, |
| 29 | init_values: float = 1e-5): |
| 30 | super().__init__() |
| 31 | self.inplace = inplace |
| 32 | self.weight = nn.Parameter(torch.ones(dim) * init_values) |
| 33 | |
| 34 | def forward(self, x): |
| 35 | if self.inplace: |
| 36 | return x.mul_(self.weight.view(-1, 1, 1)) |
| 37 | else: |
| 38 | return x * self.weight.view(-1, 1, 1) |
| 39 | |
| 40 | class TransformerStage(nn.Module): |
| 41 |