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