| 41 | |
| 42 | |
| 43 | class BertMLP(nn.Module): |
| 44 | def __init__(self, hidden_size: int, intermediate_size: int, dropout: float = 0.0): |
| 45 | super().__init__() |
| 46 | self.dense1 = nn.Linear(hidden_size, intermediate_size) |
| 47 | self.dense2 = nn.Linear(intermediate_size, hidden_size) |
| 48 | self.gelu = nn.GELU() |
| 49 | self.dropout = nn.Dropout(dropout) |
| 50 | |
| 51 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 52 | return self.dropout(self.dense2(self.gelu(self.dense1(x)))) |
| 53 | |
| 54 | |
| 55 | class BertLayer(nn.Module): |