Cannot save a Custom Variational Autoencoder Model built in Keras with tensor flow. Please help save the model

Hello to everyone reading,

Keras 2.10.0

I am trying to save my Variational Autoencoder built in Keras and Tensorflow once it is trained. I am not being able to do so. What can I do to save the model?

Attached is the Colab link to my running base model.
[Google Colab]

The error message it creates is here

ValueError: Model <__main__.VAE object at 0x00000278F0A96E60> cannot be saved either because the input shape is not available or because the forward pass of the model is not defined.To define a forward pass, please override `Model.call()`. To specify an input shape, either call `build(input_shape)` directly, or call the model on actual data using `Model()`, `Model.fit()`, or `Model.predict()`. If you have a custom training step, please make sure to invoke the forward pass in train step through `Model.__call__`, i.e. `model(inputs)`, as opposed to `model.call()`.

Furthermore I add the main lines of code here if you do not wish to open Colab.
Any help would be greatly appreciated.

import numpy as np
from tensorflow import keras
from keras.layers import Input, Dense, Lambda
import tensorflow as tf
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.datasets import make_classification
from keras import backend as K
import matplotlib.pyplot as plt
import numpy as np



# Creating fake data to replicate
# Define the desired number of samples for each class
class_samples = [2000, 3000, 2500, 1500]  # Adjust these numbers as needed

# Calculate weights based on the desired number of samples
class_weights = [num_samples / sum(class_samples) for num_samples in class_samples]

# Generate a synthetic dataset with different numbers of samples for each class
X, y = make_classification(
    n_samples=sum(class_samples),
    n_features=4,
    n_informative=4,
    n_redundant=0,
    n_classes=4,
    weights=class_weights,
    random_state=42,
)
columns = ["Feature_1", "Feature_2", "Feature_3", "Feature_4"]
synthetic_df = pd.DataFrame(data=X, columns=columns)
for column in synthetic_df:
    std = np.std(synthetic_df[column])
    mean = np.mean(synthetic_df[column])
    synthetic_df[column] = synthetic_df[column]-mean
    synthetic_df[column] = synthetic_df[column]/std
synthetic_df["target"] = y
synthetic_array =synthetic_df.values

# Defining the sampling layer that is also the call
class Sampling(keras.layers.Layer):

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

# Encoder
latent_dim = 10
encoder_inputs = Input(shape=(5), name="input_layer")

n_x = 50
x = Dense(n_x, activation="relu", name="h1")(encoder_inputs)

# Split x into two halves
half_size = n_x // 2
x3_first_half = Lambda(lambda x: x[:, :half_size], name="select_z_mean")(x)
x3_second_half = Lambda(lambda x: x[:, half_size:], name="select_z_var")(x)

z_mean = Dense(latent_dim, name="z_mean")(x3_first_half)
z_log_var = Dense(latent_dim, name="z_log_var")(x3_second_half)
z = Sampling(name="Sampling")([z_mean, z_log_var])

encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")
# Decoder
# Decoder
latent_inputs = keras.Input(shape=(latent_dim,))
n_x = 30

x = Dense(n_x, activation="relu", name="h4")(latent_inputs)

cont_decoder_outputs = Dense(4, activation="linear", name="cont_decoder_output")(x)
class_decoder_output = Dense(4, activation="softmax", name="classification_output")(x)

decoder = keras.Model(latent_inputs, [cont_decoder_outputs, class_decoder_output], name="decoder")
decoder.summary()
# Custom VAE
class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")

    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
        ]

    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            reconstruction_cont, reconstruction_class = self.decoder(z)

            data_cont = data[
                :, :4
            ]  # Assuming the first 4 columns are for continuous variables
            data_class = data[:, 4:]  # Assuming the last column is for classification

            # Reconstruction loss for continuous outputs
            reconstruction_loss_cont = keras.losses.mean_squared_error(
                data_cont, reconstruction_cont
            )

            # Reconstruction loss for classification output
            reconstruction_loss_class = keras.losses.sparse_categorical_crossentropy(
                data_class, reconstruction_class
            )

            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            kl_loss = tf.reduce_mean(kl_loss, axis=1)

            # Combine losses
            total_loss = (
                0.8*reconstruction_loss_cont + 0.2*reconstruction_loss_class + kl_loss
            )

        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)

        return {
            "loss": self.total_loss_tracker.result(),
        }

    def get_config(self):
        config = super(VAE, self).get_config()
        config.update(
            {
                "encoder": self.encoder,
                "decoder": self.decoder,
            }
        )
        return config

# Compile and Train
vae = VAE(encoder, decoder)
vae.compile(optimizer=keras.optimizers.Adam())
vae.fit(
    synthetic_array,
    epochs=1,
    batch_size=16,
)

# Save in TensorFlow SavedModel format
vae.save("path_to_save_model_tf")

Thank you very much.