unable to load a model with Keras

I wanted to create a speech recognition application and I was following this tutorial : Transformer ASR.

After training it for 1 epoch, I saved the model and it saved successfully:

model.save("model_1_epoch.keras")

But when I tried to load the model again:

custom_objects = {
    'TokenEmbedding': TokenEmbedding,
    'SpeechFeatureEmbedding': SpeechFeatureEmbedding,
    'TransformerEncoder': TransformerEncoder,
    'TransformerDecoder': TransformerDecoder,
    'Transformer': Transformer,
    'CustomSchedule': CustomSchedule
}
model = keras.models.load_model("model_1_epoch.keras", custom_objects=custom_objects, compile=False)

I get the error:

ValueError: A total of 51 objects could not be loaded. Example error message for object <Dense name=dense_65, built=True>:

Layer 'dense_65' expected 2 variables, but received 0 variables during loading. Expected: ['kernel', 'bias']

List of objects that could not be loaded:
[<Dense name=dense_65, built=True>, <Embedding name=embedding_10, built=True>, <Embedding name=embedding_11, built=True>, <EinsumDense name=key, built=True>, <EinsumDense name=attention_output, built=True>, <EinsumDense name=query, built=True>, <EinsumDense name=value, built=True>, <Dense name=dense_63, built=True>, <Dense name=dense_64, built=True>, <LayerNormalization name=layer_normalization_63, built=True>, <LayerNormalization name=layer_normalization_64, built=True>, <LayerNormalization name=layer_normalization_65 .......

This is how I am registering the custom objects:

@keras.saving.register_keras_serializable(package="SpeechFeatureEmbedding")
class SpeechFeatureEmbedding(layers.Layer):
    def __init__(self, num_hid=64, maxlen=100):
        super().__init__()
        self.num_hid = num_hid
        self.maxlen = maxlen
        self.conv1 = keras.layers.Conv1D(
            num_hid, 11, strides=2, padding="same", activation="relu"
        )
        self.conv2 = keras.layers.Conv1D(
            num_hid, 11, strides=2, padding="same", activation="relu"
        )
        self.conv3 = keras.layers.Conv1D(
            num_hid, 11, strides=2, padding="same", activation="relu"
        )
    def get_config(self):
      return{
          "num_hid": self.num_hid,
          "maxlen": self.maxlen,
      }

    def call(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return self.conv3(x)

I have used this tutorial: Keras load and save model. Especially the Registering custom objects (preferred) section as well as the passing custom object and using a custom object scope section.

But it still gives the same error no what matter what I try when I try to load the model.

I am running the code on Google Colab.

I have also tried loading without the compile argument and I have also tried saving the model as .h5 file

model.save("model.h5")

But that gives an error as well.

This is how I am creating the model:

model = Transformer(
    num_hid=200,
    num_head=2,
    num_feed_forward=400,
    target_maxlen=max_target_len,
    num_layers_enc=4,
    num_layers_dec=1,
    num_classes=34,
)

and this is how I compile and train it:

optimizer = keras.optimizers.Adam(learning_rate)
model.compile(optimizer=optimizer, loss=loss_fn)

history = model.fit(ds, validation_data=val_ds, callbacks=[display_cb], epochs=1)

using the decorator:

@keras.saving.register_keras_serializable(package="name")

and returning init params in get_config, I had to rerun the notebook and retrain the model.

Now I can successfully load the model using:

model2 = keras.models.load_model("model_1_epoch.keras", custom_objects=custom_objects, compile=False)