def METHOD_NAME(self, n, device):
    if self.mask is not None and self.mask.shape[-1] >= n:
        return self.mask[:n, :n]
    mask = torch.ones((n, n), device=device, dtype=torch.bool).triu(1)
    self.register_buffer("mask", mask, persistent=False)
    return mask