如何在TensorFlow layer(GRU)中添加打印操作？

import tensorflow as tf

# 自定义回调函数
class PrintCallback(tf.keras.callbacks.Callback):
    def on_train_batch_end(self, batch, logs=None):
        # 在每个训练批次结束时执行打印操作
        print('Batch:', batch, 'Loss:', logs['loss'])

# 创建GRU模型
model = tf.keras.Sequential([
    tf.keras.layers.GRU(64, return_sequences=True),
    tf.keras.layers.GRU(64),
    tf.keras.layers.Dense(10)
])

# 编译模型
model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True))

# 加载数据并训练模型
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train / 255.0
model.fit(x_train, y_train, epochs=10, callbacks=[PrintCallback()])