import torch


def collate_fn(data):
    """
    Creates a mini-batch, overriding default_collate function
    in order to provide batches with input sorted by length.
    """
    # sort such that input line lengths are in decreasing order
    # requirement for using torch.nn.utils.rnn.pack_packed_sequence
    data.sort(key=lambda x: x[-1], reverse=True)

    # group input sentences, target sentences, and lengths together
    inputs, targets, lens = zip(*data)
    assert len(inputs) == len(targets)

    inputs = torch.stack(inputs)
    targets = torch.stack(targets)

    return inputs, targets, list(lens)