import torch


def collate_fn(data):
    """Construct a bacth by padding the sequence to the size of the longest.

    Args:
        data (tuple): tensors

    Returns:
        tuple: padded tensors
    """
    # Construct a bacth by padding the sequence to the size of the longest
    size = [len(_x) for _x in list(zip(*data))[0]]
    pad_data = [torch.zeros(len(size), max(size)) for _ in zip(*data)]
    pad_mask = torch.ones(len(size), max(size))

    for i, _data in enumerate(data):
        end = size[i]
        pad_mask[i, :end] = 0
        for j, d in enumerate(_data):
            pad_data[j][i, :end] = d

    return [d.type(torch.int64)
            for d in pad_data] + [pad_mask.type(torch.bool)]

    # not yet supported by yapf and black formatter (allowed in Python 3.8)
    # return *[d.type(torch.int64) for d in pad_data],
    #  pad_mask.type(torch.bool)