import torch


def generate_mask(input_tensor: torch.Tensor, sequence_lengths: torch.LongTensor) -> torch.Tensor:
    """
    Generate a mask for masked 1d convolution.

    :param input_tensor: an input tensor for convolution (bsz x features x seqlen)
    :param sequence_lengths: length of sequences (bsz,)
    :return:
    """
    assert input_tensor.size()[0] == sequence_lengths.size()[0], \
        f"Batch size {input_tensor.size()[0]} != number of provided lengths {sequence_lengths.size()[0]}."

    mask = torch.ones_like(input_tensor, dtype = torch.bool)
    for i, length in enumerate(sequence_lengths):
        mask[i][:, length:] = False

    return mask