| 390 | class TransformerMLP(nn.Module): |
| 391 | |
| 392 | def __init__(self, channels, expansion, drop): |
| 393 | |
| 394 | super().__init__() |
| 395 | |
| 396 | self.dim1 = channels |
| 397 | self.dim2 = channels * expansion |
| 398 | self.chunk = nn.Sequential() |
| 399 | self.chunk.add_module('linear1', nn.Linear(self.dim1, self.dim2)) |
| 400 | self.chunk.add_module('act', nn.GELU()) |
| 401 | self.chunk.add_module('drop1', nn.Dropout(drop, inplace=True)) |
| 402 | self.chunk.add_module('linear2', nn.Linear(self.dim2, self.dim1)) |
| 403 | self.chunk.add_module('drop2', nn.Dropout(drop, inplace=True)) |
| 404 | |
| 405 | def forward(self, x): |
| 406 | |