scripts/scripts_utils/sample_batch_utils.py

def repeat_and_reshape(tensor, repeat_n):
    n_scenes = tensor.shape[0]
    return (
        tensor.view(tensor.shape[0], 1, *tensor.shape[1:])
        .repeat(1, repeat_n, *([1] * (tensor.ndim - 1)))
        .view(n_scenes * repeat_n, *tensor.shape[1:])
    )


def repeat_and_reshape_all(tensor_list, repeat_n):
    out = []
    for tensor in tensor_list:
        out.append(repeat_and_reshape(tensor, repeat_n))
    return out