(self, in_dim, out_dim, has_pos_emb=False)
| 477 | |
| 478 | class MLP(torch.nn.Module): |
| 479 | def __init__(self, in_dim, out_dim, has_pos_emb=False): |
| 480 | super().__init__() |
| 481 | self.proj = torch.nn.Sequential( |
| 482 | nn.LayerNorm(in_dim), |
| 483 | nn.Linear(in_dim, in_dim), |
| 484 | nn.GELU(), |
| 485 | nn.Linear(in_dim, out_dim), |
| 486 | nn.LayerNorm(out_dim) |
| 487 | ) |
| 488 | self.has_pos_emb = has_pos_emb |
| 489 | if has_pos_emb: |
| 490 | self.emb_pos = torch.nn.Parameter(torch.zeros((1, 514, 1280))) |
| 491 | |
| 492 | def forward(self, x): |
| 493 | if self.has_pos_emb: |