使用tensorflow打印自定义train_step函数中的值

import tensorflow as tf

@tf.function
def train_step(inputs, labels):
    with tf.GradientTape() as tape:
        predictions = model(inputs)
        loss = loss_function(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    
    # 打印损失值
    print("Batch Loss:", loss.numpy())

for inputs, labels in train_dataset:
    train_step(inputs, labels)