Method __init__

tests/unit/runtime/zero/test_zero_leaf_module.py:99–104

(self, hidden_dim)

class MLPBlock(nn.Module):

    def __init__(self, hidden_dim):
        super(MLPBlock, self).__init__()
        self.gate_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.up_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.act_fn = nn.GELU()

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))