Customizing fit() function in keras leads evaluate() function returning empty list

Question

I am training a VAE using the Keras tutorial: https://keras.io/examples/generative/vae/. This involves creating a VAE class and specifying a custom training procedure, which is described more here: https://keras.io/guides/customizing_what_happens_in_fit/. After creating an encoder and decoder as specified in the tutorial and training a model, I create the VAE model through the following:

vae = VAE(encoder, decoder) 
vae.compile(optimizer=keras.optimizers.Adam())

vae.fit(x=x_train, y=None, 
        epochs=epochs, 
        batch_size=batch_size, 
        verbose=False,
        validation_data=(x_test, None))

I want to evaluate the model on separate datasets after training (since I have multiple evaluation datasets, I am not using it as validation_data). However, when I try to run vae.evaluate(data), it returns an empty list [].

Note: I am able to obtain training and validation metrics easily through vae.history.history, but the issue is when I try to evaluate after training. But when I try to return the metrics vae.metrics it also returns an empty list. How do I get model.evaluate to work with a custom training procedure that returns a dict of loss metrics? Do I need to define something custom for evaluate?

Here is how the VAE class is defined. More details can be found in the tutorial above.

class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def train_step(self, data):
        if isinstance(data, tuple):
            data = data[0]
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = encoder(data)
            reconstruction = decoder(z)
            reconstruction_loss = tf.reduce_mean(
                keras.losses.binary_crossentropy(data, reconstruction)
            )
            reconstruction_loss *= 28 * 28
            kl_loss = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
            kl_loss = tf.reduce_mean(kl_loss)
            kl_loss *= -0.5
            total_loss = reconstruction_loss + kl_loss
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        return {
            "loss": total_loss,
            "reconstruction_loss": reconstruction_loss,
            "kl_loss": kl_loss,
        }