在使用customize train_step()的Keras中，如何在使用customize test_step()评估模型时设置'training=False‘？

@tf.function
def train_step(inputs, labels, model, optimizer, loss_object):
    with tf.GradientTape() as tape:
        predictions = model(inputs, training=True)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

@tf.function
def test_step(inputs, labels, model, loss_object):
    predictions = model(inputs, training=False)
    loss = loss_object(labels, predictions)
    return loss

# 在使用customize train_step()和test_step()的训练循环中，我们可以这样设置'training=False'：

for epoch in range(num_epochs):
    for inputs, labels in train_dataset:
        train_loss = train_step(inputs, labels, model, optimizer, loss_object)
    
    for test_inputs, test_labels in test_dataset:
        test_loss = test_step(test_inputs, test_labels, model, loss_object)