| 201 | return h_.sequential(self.to_out) |
| 202 | |
| 203 | class GEGLU: |
| 204 | def __init__(self, dim_in, dim_out): |
| 205 | self.proj = Linear(dim_in, dim_out * 2) |
| 206 | self.dim_out = dim_out |
| 207 | |
| 208 | def __call__(self, x): |
| 209 | x, gate = self.proj(x).chunk(2, dim=-1) |
| 210 | return x * gate.gelu() |
| 211 | |
| 212 | class FeedForward: |
| 213 | def __init__(self, dim, mult=4): |