import torch


def multitensor_collate_fxn(x):
    """
    Takes a list of BATCHSIZE lists of tensors of length D.
    Returns a list of length D of batched tensors.
    """
    num_tensors_to_batch = len(x[0])

    regroup_tensors = [[] for i in range(num_tensors_to_batch)]
    for t_list in x:
        for i, t in enumerate(t_list):
            regroup_tensors[i].append(t.unsqueeze(0))
    batched_tensors = [torch.cat(tl) for tl in regroup_tensors]
    return batched_tensors