import torch


def pack_tensors(tensors, use_cuda=False):
    """
    Packs a list of tensors into one 1-dimensional tensor.

    Args:
        tensors (list[torch.Tensor]): The tensors to pack
        use_cuda (bool): Whether the resulting tensor should be on cuda

    Returns:
        (torch.Tensor, list[int], list[(int, int)]):
            The flattened tensors, the list start indices of each packed tensor,
            and the original shape of each tensor.

            Those values are used to then unpack the tensor
    """
    indices = [0]
    for tensor in tensors:
        new_end = indices[-1] + tensor.nelement()
        indices.append(new_end)

    tensor_sizes = [t.size() for t in tensors]

    vec = torch.empty(
        indices[-1],
        device=tensors[0].device if tensors[0].is_cuda and use_cuda else "cpu",
        dtype=tensors[0].dtype,
    )

    for tensor, start_idx, end_idx in zip(tensors, indices[:-1], indices[1:]):
        vec[start_idx:end_idx] = tensor.data.view(-1)

    return vec, indices, tensor_sizes