| 184 | |
| 185 | |
| 186 | class Residual(Module): |
| 187 | @beartype |
| 188 | def __init__(self, fn: Module): |
| 189 | super().__init__() |
| 190 | self.fn = fn |
| 191 | |
| 192 | def forward(self, x, **kwargs): |
| 193 | return self.fn(x, **kwargs) + x |
| 194 | |
| 195 | |
| 196 | # for a bunch of tensor operations to change tensor to (batch, time, feature dimension) and back |