如何在VGG16的Tensorflow中采用渐变

import tensorflow as tf
from tensorflow.keras.applications import VGG16

vgg16 = VGG16(weights='imagenet', include_top=False)

input_tensor = tf.keras.Input(shape=(224, 224, 3))
output = vgg16(input_tensor)

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
optimizer = tf.keras.optimizers.SGD(learning_rate=0.001)

@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        predictions = vgg16(images, training=True)
        loss = loss_fn(labels, predictions)
    gradients = tape.gradient(loss, vgg16.trainable_variables)
    optimizer.apply_gradients(zip(gradients, vgg16.trainable_variables))

for images, labels in train_dataset:
    train_step(images, labels)