在tf.keras中使用Horovod时，如何从检查点恢复？

import tensorflow as tf
import horovod.tensorflow.keras as hvd

hvd.init()

config = tf.ConfigProto()
config.gpu_options.visible_device_list = str(hvd.local_rank())
tf.keras.backend.set_session(tf.Session(config=config))

model = tf.keras.models.Sequential()
# 添加模型层

optimizer = tf.keras.optimizers.Adam(0.001 * hvd.size())
optimizer = hvd.DistributedOptimizer(optimizer)
model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['accuracy'])

checkpoint_dir = './checkpoints'
if hvd.rank() == 0:
    os.makedirs(checkpoint_dir, exist_ok=True)
checkpoint_file = os.path.join(checkpoint_dir, 'checkpoint.h5')
checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(checkpoint_file, save_weights_only=True)

if os.path.exists(checkpoint_file):
    model.load_weights(checkpoint_file)

model.fit(x_train, y_train, callbacks=[checkpoint_callback], ...)