| 585 | |
| 586 | |
| 587 | class FeedForward(nn.Module): |
| 588 | |
| 589 | def __init__(self, dim, dim_out=None, mult=4, glu=False): |
| 590 | super().__init__() |
| 591 | inner_dim = int(dim * mult) |
| 592 | self.net = nn.Sequential( |
| 593 | GEGLU(dim_in=dim, dim_out=inner_dim), nn.Identity(), |
| 594 | nn.Conv2d(inner_dim, |
| 595 | dim_out if dim_out is not None else dim, |
| 596 | kernel_size=1)) |
| 597 | |
| 598 | def forward(self, hidden_states): |
| 599 | return self.net(hidden_states) |
| 600 | |
| 601 | |
| 602 | class GEGLU(nn.Module): |