有没有办法在Keras框架中使用global_step？

import tensorflow as tf
from tensorflow.keras import layers, models

# 创建一个简单的模型
model = models.Sequential([
    layers.Dense(64, activation='relu', input_shape=(784,)),
    layers.Dense(10, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# 创建一个Checkpoint对象
checkpoint = tf.train.Checkpoint(model=model, optimizer=tf.keras.optimizers.Adam())

# 创建一个CheckpointManager对象
checkpoint_manager = tf.train.CheckpointManager(checkpoint, directory='./checkpoints', max_to_keep=3)

# 定义global_step
global_step = tf.Variable(0, trainable=False, dtype=tf.int64)

# 自定义训练循环
for epoch in range(epochs):
    for batch, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        with tf.GradientTape() as tape:
            logits = model(x_batch_train, training=True)
            loss_value = loss_fn(y_batch_train, logits)
        
        grads = tape.gradient(loss_value, model.trainable_weights)
        optimizer.apply_gradients(zip(grads, model.trainable_weights))
        
        # 更新global_step
        global_step.assign_add(1)
        
        # 保存checkpoint
        if global_step % 100 == 0:
            checkpoint_manager.save(checkpoint_number=global_step)
    
    print(f'Epoch {epoch + 1}, Loss: {loss_value.numpy()}')