如何在训练运行之间的中间层中更新每个时期的参数？(tensorflow急切执行)

import tensorflow as tf

# 定义模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

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

# 定义优化器
optimizer = tf.keras.optimizers.SGD(learning_rate=0.001)

# 定义训练循环
for epoch in range(num_epochs):
    # 在每个时期开始前，初始化梯度
    grads = tf.GradientTape()

    # 前向传播计算预测结果
    logits = model(x_train)

    # 计算损失函数
    loss_value = loss_fn(y_train, logits)

    # 计算梯度
    grads = tape.gradient(loss_value, model.trainable_variables)

    # 更新参数
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    # 打印每个时期的损失值
    print(f"Epoch {epoch+1}: Loss = {loss_value.numpy()}")

# 使用训练好的模型进行预测
predictions = model(x_test)