如何在tensorflow中实现iter_size like caffe

import tensorflow as tf

# 定义数据集
dataset = tf.data.Dataset.from_tensor_slices((features, labels))
dataset = dataset.shuffle(buffer_size=1000).batch(batch_size)

# 定义模型
model = create_model()

# 定义优化器
optimizer = tf.keras.optimizers.Adam()

# 定义损失函数
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()

# 定义准确率指标
accuracy_metric = tf.keras.metrics.SparseCategoricalAccuracy()

# 定义迭代次数
iter_size = 4

# 迭代训练
for epoch in range(num_epochs):
    for batch in dataset:
        # 梯度累积
        for _ in range(iter_size):
            with tf.GradientTape() as tape:
                logits = model(batch[0], training=True)
                loss_value = loss_fn(batch[1], logits)
            grads = tape.gradient(loss_value, model.trainable_variables)
            optimizer.apply_gradients(zip(grads, model.trainable_variables))
        
        # 更新指标
        accuracy_metric(batch[1], logits)
        
    # 打印训练结果
    print("Epoch {}: Loss = {}, Accuracy = {}".format(epoch, loss_value, accuracy_metric.result()))
    accuracy_metric.reset_states()