(self, module, scales)
| 30 | |
| 31 | class ScaledActivation(nn.Module): |
| 32 | def __init__(self, module, scales): |
| 33 | super().__init__() |
| 34 | self.act = module |
| 35 | self.scales = nn.Parameter(scales.data) |
| 36 | |
| 37 | def forward(self, x): |
| 38 | return self.act(x) / self.scales.view(1, 1, -1).to(x.device) |