| 25 | |
| 26 | |
| 27 | class PreNorm(nn.Module): |
| 28 | def __init__(self, dim, fn): |
| 29 | super().__init__() |
| 30 | self.norm = nn.LayerNorm(dim) |
| 31 | self.fn = fn |
| 32 | |
| 33 | def forward(self, x, **kwargs): |
| 34 | return self.fn(self.norm(x), **kwargs) |
| 35 | |
| 36 | |
| 37 | class FeedForward(nn.Module): |