| 92 | |
| 93 | |
| 94 | class SwiGLU(nn.Module): |
| 95 | |
| 96 | def __init__(self, dim, mid_dim): |
| 97 | super().__init__() |
| 98 | self.dim = dim |
| 99 | self.mid_dim = mid_dim |
| 100 | |
| 101 | # layers |
| 102 | self.fc1 = nn.Linear(dim, mid_dim) |
| 103 | self.fc2 = nn.Linear(dim, mid_dim) |
| 104 | self.fc3 = nn.Linear(mid_dim, dim) |
| 105 | |
| 106 | def forward(self, x): |
| 107 | x = F.silu(self.fc1(x)) * self.fc2(x) |
| 108 | x = self.fc3(x) |
| 109 | return x |
| 110 | |
| 111 | |
| 112 | class AttentionBlock(nn.Module): |