| 210 | return x * gate.gelu() |
| 211 | |
| 212 | class FeedForward: |
| 213 | def __init__(self, dim, mult=4): |
| 214 | self.net = [ |
| 215 | GEGLU(dim, dim*mult), |
| 216 | lambda x: x, # needed for weights loading code to work |
| 217 | Linear(dim*mult, dim) |
| 218 | ] |
| 219 | |
| 220 | def __call__(self, x): |
| 221 | return x.sequential(self.net) |
| 222 | |
| 223 | class BasicTransformerBlock: |
| 224 | def __init__(self, dim, context_dim, n_heads, d_head): |