| 600 | |
| 601 | |
| 602 | class GEGLU(nn.Module): |
| 603 | |
| 604 | def __init__(self, dim_in, dim_out): |
| 605 | super().__init__() |
| 606 | self.proj = nn.Conv2d(dim_in, dim_out * 2, kernel_size=1) |
| 607 | |
| 608 | def forward(self, hidden_states): |
| 609 | hidden_states, gate = self.proj(hidden_states).chunk(2, dim=1) |
| 610 | return hidden_states * F.gelu(gate) |
| 611 | |
| 612 | |
| 613 | def get_activation(act_fn): |