Получение AttributeError: объект 'Tensor' не имеет атрибута 'log_prob' при сохранении модели тензорного потока

Question

Я пытаюсь сохранить вариационный автокодер, построенный с использованием вероятностей TensorFlow и Tensorflow. Моя цель - обслуживать модель с помощью обслуживания TensorFlow. Я открыт для обслуживания этой модели другим способом.

Вот моя модель:

def dense_layers(sizes):
return tfk.Sequential([tfkl.Dense(size, activation=tf.nn.leaky_relu) for size in sizes])

original_dim = 30
input_shape = (30,)
intermediary_dims = [20, 10, 8]
latent_dim = 2


prior = tfd.MultivariateNormalDiag(
        loc=tf.zeros([latent_dim]),
        scale_identity_multiplier=1.0)

encoder = tfk.Sequential([
    tfkl.InputLayer(input_shape=input_shape, name='encoder_input'),
    dense_layers(intermediary_dims),
    tfkl.Dense(tfpl.MultivariateNormalTriL.params_size(latent_dim), activation=None),
    tfpl.MultivariateNormalTriL(latent_dim, 
                           activity_regularizer=tfpl.KLDivergenceRegularizer(prior)),
], name='encoder')

encoder.summary()

decoder = tfk.Sequential([
    tfkl.InputLayer(input_shape=[latent_dim]),
    dense_layers(reversed(intermediary_dims)),
    tfkl.Dense(tfpl.IndependentNormal.params_size(original_dim), activation=None),
    tfpl.IndependentNormal(original_dim),
], name='decoder')

decoder.summary()

vae = tfk.Model(inputs=encoder.inputs,
                outputs=decoder(encoder.outputs[0]),
                name='vae_mlp')

negloglik = lambda x, rv_x: -rv_x.log_prob(x)

vae.compile(optimizer=tf.keras.optimizers.RMSprop(), 
            loss=negloglik)

vae.summary()

Вот как я пытаюсь сохранить модель:

tf.keras.models.save_model(
                vae,
                "/opt/notebooks/saved/vae/1",
                overwrite=True,
                include_optimizer=True,
                save_format=None,
                signatures=None,
                options=None
            )