Function train_step

tests/performance_test/vgg/tf2-autograph.py:39–48

(x_batch, y_batch)

# training function
@tf.function
def train_step(x_batch, y_batch):
    # forward + backward
    with tf.GradientTape() as tape:
        ## compute outputs
        _logits = vgg(x_batch, training=True)
        ## compute loss and update model
        _loss = loss_object(y_batch, _logits)

    grad = tape.gradient(_loss, train_weights)
    optimizer.apply_gradients(zip(grad, train_weights))


# begin training