| 209 | |
| 210 | |
| 211 | class FalconMLP(nn.Module): |
| 212 | def __init__( |
| 213 | self, |
| 214 | dim: int, |
| 215 | ): |
| 216 | super().__init__() |
| 217 | self.dense_h_to_4h = nn.Linear(dim, 4 * dim, bias=False) |
| 218 | self.act = nn.GELU() |
| 219 | self.dense_4h_to_h = nn.Linear(4 * dim, dim, bias=False) |
| 220 | |
| 221 | def forward(self, x): |
| 222 | x = self.act(self.dense_h_to_4h(x)) |
| 223 | x = self.dense_4h_to_h(x) |
| 224 | return x |
| 225 | |
| 226 | |
| 227 | class TransformerBlock(nn.Module): |