import torch


def pad_batch(items, pad_value=0):
    """Pad tensors in list to equal length (in the first dim)

    :param items:
    :param pad_value:
    :return: padded_items, orig_lens
    """
    max_len = len(max(items, key=lambda x: len(x)))
    zeros = (2*torch.as_tensor(items[0]).ndim -1) * [pad_value]
    padded_items = torch.stack([torch.nn.functional.pad(torch.as_tensor(x), pad= zeros + [max_len - len(x)], value=pad_value)
                          for x in items])
    orig_lens = [len(xx) for xx in items]
    return padded_items, orig_lens