| 584 | |
| 585 | class FeedForward(nn.Module): |
| 586 | def __init__( |
| 587 | self, |
| 588 | dim: int, |
| 589 | inner_dim: Optional[int] = None, |
| 590 | dim_out: Optional[int] = None, |
| 591 | mult: int = 4, |
| 592 | bias: bool = False, |
| 593 | ): |
| 594 | super().__init__() |
| 595 | inner_dim = dim*mult if inner_dim is None else inner_dim |
| 596 | dim_out = dim if dim_out is None else dim_out |
| 597 | self.net = nn.ModuleList([ |
| 598 | GELU(dim, inner_dim, approximate="tanh", bias=bias), |
| 599 | nn.Identity(), |
| 600 | nn.Linear(inner_dim, dim_out, bias=bias) |
| 601 | ]) |
| 602 | |
| 603 | |
| 604 | def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: |