| 153 | |
| 154 | |
| 155 | class LinearGEGLU(nn.Linear): |
| 156 | def __init__(self, in_features, out_features, bias=True): |
| 157 | super().__init__(in_features, out_features * 2, bias=bias) |
| 158 | self.out_features = out_features |
| 159 | |
| 160 | def forward(self, x): |
| 161 | flops.op(flops.op_linear, x.shape, self.weight.shape) |
| 162 | return linear_geglu(x, self.weight, self.bias) |
| 163 | |
| 164 | |
| 165 | class RMSNorm(nn.Module): |