import torch


def create_length_mask(data, lengths):
    """
    Create lengths mask for data along one dimension.
    """
    n_sequences, max_length, _ = data.shape
    lengths_mask = torch.zeros(n_sequences, max_length)
    for i, length in enumerate(lengths):
        lengths_mask[i, :length + 1] = 1
    return lengths_mask