Method forward

transformer.py:164–166

(self, x)

        self.down_proj = Linear(expand_dim, dim)

    def forward(self, x):
        gate, up = self.gateup_proj(x).chunk(2, dim=-1)
        return self.down_proj(up * F.silu(gate))

class GQA(nn.Module):
    def __init__(self,