如何使用Tensorflow 2.1保存检查点的平均权重？

import tensorflow as tf
import numpy as np

# 定义一个简单的模型
class SimpleModel(tf.keras.Model):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.dense = tf.keras.layers.Dense(10)

    def call(self, inputs):
        return self.dense(inputs)

# 创建模型和优化器
model = SimpleModel()
optimizer = tf.keras.optimizers.Adam()

# 创建检查点管理器
checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
checkpoint_manager = tf.train.CheckpointManager(checkpoint, './checkpoints', max_to_keep=5)

# 训练模型并保存检查点
def train_step(inputs, targets):
    with tf.GradientTape() as tape:
        predictions = model(inputs)
        loss = tf.keras.losses.mean_squared_error(targets, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

# 模拟训练过程
for epoch in range(10):
    inputs = np.random.rand(32, 5).astype(np.float32)
    targets = np.random.rand(32, 10).astype(np.float32)
    loss = train_step(inputs, targets)
    print(f"Epoch {epoch}, Loss: {loss.numpy().mean()}")
    checkpoint_manager.save()

# 计算并保存平均权重
def average_checkpoints(checkpoint_manager):
    num_checkpoints = len(checkpoint_manager.checkpoints)
    if num_checkpoints == 0:
        return

    # 初始化平均权重
    avg_weights = [tf.zeros_like(var) for var in model.trainable_variables]

    # 累加每个检查点的权重
    for checkpoint_path in checkpoint_manager.checkpoints:
        checkpoint.restore(checkpoint_path).expect_partial()
        for i, var in enumerate(model.trainable_variables):
            avg_weights[i] += var / num_checkpoints

    # 将平均权重赋值给模型
    for i, var in enumerate(model.trainable_variables):
        var.assign(avg_weights[i])

    # 保存平均权重
    avg_checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
    avg_checkpoint_manager = tf.train.CheckpointManager(avg_checkpoint, './avg_checkpoints', max_to_keep=1)
    avg_checkpoint_manager.save()

# 计算并保存平均权重
average_checkpoints(checkpoint_manager)