Tensorflow 2.0 Keras BatchNorm:如何在自定义培训中更新在线参数？

import tensorflow as tf
from tensorflow.keras.layers import BatchNormalization

class CustomModel(tf.keras.Model):
    def __init__(self):
        super(CustomModel, self).__init__()
        self.batch_norm = BatchNormalization()

    def call(self, inputs, training=False):
        x = self.batch_norm(inputs, training=training)
        # 添加其他层和逻辑
        return x

loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()

model = CustomModel()

def compute_loss(model, x, y):
    logits = model(x, training=True)
    loss = loss_object(y, logits)
    return loss

@tf.function
def train_step(model, x, y):
    with tf.GradientTape() as tape:
        loss = compute_loss(model, x, y)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

# 假设有训练数据x和标签y
train_step(model, x, y)