Method forward

transformer.py:319–320

(self, x)

        nn.init.trunc_normal_(self.output.weight, std=std, a=-3 * std, b=3 * std)

    def forward(self, x):
        return self.output(self.norm(x))
    
@si_module
class Stack(nn.Module):