import torch


def mask_forward(sz: int, diagonal: int = 1):# -> torch.Tensor:
    """
    Generate a mask that ignores future words. Each (i, j)-entry will be True if j >= i + diagonal

    Args:
        sz (int): Length of the sequence.
        diagonal (int): Amount of shift for diagonal entries.
    
    Returns: 
        torch.Tensor: Mask tensor with shape [sz, sz].
    """
    return torch.ones(sz, sz, dtype=torch.bool, requires_grad=False).triu(diagonal=diagonal).contiguous()