Method forward

bitnet/bit_mamba.py:174–180

(self, x)

        self.norm = RMSNorm(config.dim)

    def forward(self, x):
        # x : (B, L, D)

        # output : (B, L, D)

        output = self.mixer(self.norm(x)) + x
        return output

    def step(self, x, cache):
        # x : (B, D)