Try this: import tensorflow as tfimport numpy as npdef wrap(dist): return tf.while_loop( cond=lambda X: tf.math.abs(X) > 0.5, body=lambda X: tf.math.subtract(X, 1.0), loop_vars=(dist))def custom_loss(y_true, y_pred): diff = tf.math.abs(y_true - y_pred) diff = tf.reshape(diff, [-1]) diff = tf.vectorized_map(wrap, [diff]) return tf.math.reduce_mean(tf.math.square(diff))y_true = np.array([[0., 1., 1.0], [0., 0., 0.]])y_pred = np.array([[1., 1., 1.0], [1., 0., 1.]])custom_loss(y_true, y_pred).numpy()