Iterate in Keras custom loss function
Assuming 2 partitions since, for index in range(0,2) and a fixed bunch size of 3 since, start_range + 3
The shape of your y_true should be (nb_samples, 2), where the first vector of 2nd dimension is the ground_truth and the other is partition_index
import numpy as np
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.keras.losses import mean_squared_error

y_true = tf.Variable(np.array([[1.5, 0],[1.2, 0],[1.3, 0],[1.6, 1],[3.0, 1],[2.25, 1]]), dtype=tf.float32)
y_pred = tf.Variable(np.array([[1.35],[1.24],[1.69],[1.55],[1.24],[1.69]]), dtype=tf.float32)

def get_two_splits(y_true):
    child_y = y_true[:, 1]
    child_y = tf.expand_dims(child_y, 1)
    y_true = y_true[:, 0]
    y_true = tf.expand_dims(y_true, 1)
    return y_true, child_y

def some_processing_on_bunches(y_all):
    y_true, y_pred = get_two_splits(y_all)
    return mean_squared_error(y_true, y_pred)

def custom_loss(y_true, y_pred):
    n = K.shape(y_true)[0]
    y_true, partitions = get_two_splits(y_true)
    partitions = tf.reshape(partitions ,[n, ])
    partitions = tf.cast(partitions, tf.int32)

    y_elements = tf.dynamic_partition(tf.concat([y_true, y_pred], 1), partitions, 2)
    losses = tf.vectorized_map(some_processing_on_bunches, tf.stack(y_elements))
    return tf.reduce_mean(losses, axis=1)

K.eval(custom_loss(y_true, y_pred))
&gt;&gt;&gt;array([0.05873336, 1.1379    ], dtype=float32)