def _predict_batch_by_batch(predict_fn, batched_inputs, non_batched_inputs=None, batch_size=BATCH_SIZE, num_outputs=1):
    """
    Splits prediction for big dataset in order to save GPU memory.
    Equivalent to predict_fn(*batched_inputs + non_batched_inputs).

    predict_fn: compiled keras predict function.
    batched_inputs:
        Inputs that we split into batches. On each iteration, we only pass one batch of this data into predict_fn.
    non_batched_inputs:
        Inputs that we do not split into batches. These inputs are the same for each call of predict_fn
    batch_size: int
        Size of each batch that we split our batched_inputs into
    num_ouputs: int, default=1
        Number of items returned on each call of predict_fn.
    """
    if non_batched_inputs is None:
        non_batched_inputs = []

    results = [[] for _ in range(num_outputs)]

    for inputs_batch in get_training_batch(batched_inputs, batch_size):
        args = list(inputs_batch) + non_batched_inputs
        cur_result = predict_fn(*args)
        if num_outputs > 1:
            for i in range(num_outputs):
                results[i].append(cur_result[i])
        else:
            results[0].append(cur_result)

    if num_outputs > 1:
        return tuple(np.concatenate(results[i]) for i in range(num_outputs))
    else:
        return np.concatenate(results[0])

    Splits prediction for big dataset in order to save GPU memory.
    Equivalent to predict_fn(*batched_inputs + non_batched_inputs).

    predict_fn: compiled keras predict function.
    batched_inputs:
        Inputs that we split into batches. On each iteration, we only pass one batch of this data into predict_fn.
    non_batched_inputs:
        Inputs that we do not split into batches. These inputs are the same for each call of predict_fn
    batch_size: int
        Size of each batch that we split our batched_inputs into
    num_ouputs: int, default=1
        Number of items returned on each call of predict_fn.
    
cakechat/dialog_model/inference/utils.py