如何在tensorflow 2.0中训练tensorflow概率双射(权重)？

import tensorflow as tf
import tensorflow_probability as tfp

model = tf.keras.Sequential([
    tfp.layers.DenseVariational(units=hidden_units, 
                                make_posterior_fn=tfp.layers.util.default_mean_field_normal_fn(),
                                make_prior_fn=tfp.layers.util.default_multivariate_normal_fn,
                                kl_weight=1/num_train_samples),
    tfp.layers.DistributionLambda(lambda t: tfp.distributions.Normal(loc=t, scale=1)),
])

def negative_log_likelihood(y_true, y_pred):
    return -y_pred.log_prob(y_true)

model.compile(optimizer=tf.keras.optimizers.Adam(), loss=negative_log_likelihood)

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train / 255.0
x_test = x_test / 255.0

dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
dataset = dataset.shuffle(buffer_size=1024).batch(batch_size)

model.fit(dataset, epochs=num_epochs)