(self, input_dim, hidden_dim, output_dim)
| 279 | |
| 280 | class alignMLP(nn.Module): |
| 281 | def __init__(self, input_dim, hidden_dim, output_dim): |
| 282 | super(alignMLP, self).__init__() |
| 283 | self.layer1 = nn.Linear(input_dim, hidden_dim) |
| 284 | self.layer2 = nn.Linear(hidden_dim, hidden_dim) |
| 285 | self.layer3 = nn.Linear(hidden_dim, hidden_dim) |
| 286 | self.layer4 = nn.Linear(hidden_dim, hidden_dim) |
| 287 | self.layer5 = nn.Linear(hidden_dim, output_dim) |
| 288 | self.norm = nn.LayerNorm(output_dim) |
| 289 | |
| 290 | def forward(self, x): |
| 291 | x = F.gelu(self.layer1(x)) |