| 18 | |
| 19 | |
| 20 | class GEGLU(torch.nn.Module): |
| 21 | |
| 22 | def __init__(self, dim_in, dim_out): |
| 23 | super().__init__() |
| 24 | self.proj = torch.nn.Linear(dim_in, dim_out * 2) |
| 25 | |
| 26 | def forward(self, hidden_states): |
| 27 | hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1) |
| 28 | return hidden_states * torch.nn.functional.gelu(gate) |
| 29 | |
| 30 | |
| 31 | class BasicTransformerBlock(torch.nn.Module): |