| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class MLP(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.fc_1 = nn.Linear(config.embed_dim, config.mlp_dim, bias=True) | |
| self.fc_2 = nn.Linear(config.mlp_dim, config.embed_dim, bias=True) | |
| self.activation = nn.GELU() | |
| self.dropout = nn.Dropout(0.1) | |
| def forward(self, x): | |
| x = self.fc_1(x) | |
| x = self.activation(x) | |
| x = self.dropout(x) | |
| x = self.fc_2(x) | |
| return x |