在自定义损失函数中使用tf.while_loop的正确方法是什么？

import tensorflow as tf

def custom_loss_function(labels, predictions):
    # 定义循环体函数
    def loop_body(i, loss):
        # 获取当前样本的标签和预测值
        label = labels[i]
        prediction = predictions[i]
        
        # 计算当前样本的损失
        sample_loss = tf.square(label - prediction)
        
        # 更新总体损失
        loss += sample_loss
        
        # 更新循环变量
        i += 1
        
        return i, loss

    # 初始化循环变量
    i = tf.constant(0)
    total_loss = tf.constant(0.0)

    # 使用tf.while_loop执行循环
    _, loss = tf.while_loop(
        lambda i, loss: tf.less(i, tf.shape(labels)[0]),  # 循环条件函数
        loop_body,  # 循环体函数
        loop_vars=[i, total_loss],  # 循环变量
        shape_invariants=[i.get_shape(), total_loss.get_shape()]  # 循环变量的形状不变性
    )

    # 计算平均损失
    loss = tf.reduce_mean(loss)
    
    return loss