(self, dim, mid_dim)
| 271 | class SwiGLU(nn.Module): |
| 272 | |
| 273 | def __init__(self, dim, mid_dim): |
| 274 | super().__init__() |
| 275 | self.dim = dim |
| 276 | self.mid_dim = mid_dim |
| 277 | |
| 278 | # layers |
| 279 | self.fc1 = nn.Linear(dim, mid_dim) |
| 280 | self.fc2 = nn.Linear(dim, mid_dim) |
| 281 | self.fc3 = nn.Linear(mid_dim, dim) |
| 282 | |
| 283 | def forward(self, x): |
| 284 | x = F.silu(self.fc1(x)) * self.fc2(x) |