import torch


def seq_mask(seq_len, max_len):
    """Create a mask for the sequences.

    :param seq_len: list or torch.LongTensor
    :param max_len: int
    :return mask: torch.LongTensor
    """
    if isinstance(seq_len, list):
        seq_len = torch.LongTensor(seq_len)
    mask = [torch.ge(seq_len, i + 1) for i in range(max_len)]
    mask = torch.stack(mask, 1)
    return mask