| 115 | |
| 116 | # classes |
| 117 | class Scale(nn.Module): |
| 118 | def __init__(self, value, fn): |
| 119 | super().__init__() |
| 120 | self.value = value |
| 121 | self.fn = fn |
| 122 | |
| 123 | def forward(self, x, **kwargs): |
| 124 | x, *rest = self.fn(x, **kwargs) |
| 125 | return (x * self.value, *rest) |
| 126 | |
| 127 | |
| 128 | class Rezero(nn.Module): |