self.register_buffer(
    "mask",
    torch.triu(
        torch.ones(context_length, context_length),
        diagonal=1,
    ),
)

        __
        __