| 306 | return x |
| 307 | |
| 308 | class FeedForward(nn.Module): |
| 309 | def __init__(self, in_dim, hidden_dim, out_chans=None, act_layer=nn.GELU, dropout=0.): |
| 310 | super().__init__() |
| 311 | |
| 312 | out_chans = out_chans or in_dim |
| 313 | hidden_dim = hidden_dim or in_dim |
| 314 | |
| 315 | self.net = nn.Sequential( |
| 316 | nn.Linear(in_dim, hidden_dim), |
| 317 | act_layer(), |
| 318 | nn.Dropout(dropout), |
| 319 | nn.Linear(hidden_dim, out_chans), |
| 320 | nn.Dropout(dropout) |
| 321 | ) |
| 322 | |
| 323 | def forward(self, x): |
| 324 | return self.net(x) |
| 325 | |
| 326 | class Attention(nn.Module): |
| 327 | def __init__(self, |