import torch


def get_mask(tensor, lengths):
    """
    Generates the masks for batches of sequences.
    Time should be the first axis.
    input:
        tensor: the tensor for which to generate the mask [N x T x ...]
        lengths: lengths of the seq. [N] 
    """
    mask = torch.zeros_like(tensor)
    for i in range(len(lengths)):
        mask[i, :lengths[i]] = 1.
    return mask