I have made a Custom Keras model for Machine Translation. After training and saving the model as tf_model.keras , during loading I am getting an error. I have added get_config() function and also added @keras.saving.register_keras_serializable() at the top of the class.

Code for the Model :

@keras.saving.register_keras_serializable()
class BidirectionalEncoderandDecoderWithAttention(keras.Model):
  def __init__(
      self,
      input_vocabulary_size = 5000,
      output_vocabulary_size=5000,
      embedding_size = 256,
      units_lstm =512,
      **kwargs):
    super().__init__(**kwargs)
    self.input_vocabulary_size = input_vocabulary_size
    self.output_vocabulary_size = output_vocabulary_size
    self.embedding_size = embedding_size
    self.units_lstm = units_lstm
    self.encoder_embedding = Embedding(self.input_vocabulary_size,self.embedding_size,mask_zero = True)
    self.encoder = Bidirectional(LSTM(self.units_lstm//2,return_sequences = True,return_state = True))
    self.decoder_embedding = Embedding(self.output_vocabulary_size,self.embedding_size,mask_zero = True)
    self.decoder = LSTM(self.units_lstm,return_sequences = True)
    self.attention = Attention()
    self.output_layer = Dense(self.output_vocabulary_size,activation = 'softmax')
  def call(self,inputs):
      encoder_inputs,decoder_inputs = inputs
      encoder_embeddings = self.encoder_embedding(encoder_inputs)
      decoder_embeddings = self.decoder_embedding(decoder_inputs)

      encoder_op,*encoder_state = self.encoder(encoder_embeddings)
      encoder_state = [
            tf.concat(encoder_state[0::2], axis=-1),
            tf.concat(encoder_state[1::2], axis=-1),
        ]
      decoder_op = self.decoder(decoder_embeddings,initial_state = encoder_state)
      attention_output = self.attention([decoder_op,encoder_op])
      output = self.output_layer(attention_output)
      return output
  def get_config(self):
        config = super().get_config()
        config.update ({
            'input_vocabulary_size':self.input_vocabulary_size,
            'output_vocabulary_size':self.output_vocabulary_size,
            'embedding_size':self.embedding_size,
            'units_lstm':self.units_lstm,
            'encoder_embedding':self.encoder_embedding,
            'encoder':self.encoder,
            'decoder_embedding' : self.decoder_embedding,
            'decoder': self.decoder,
            'attention':self.attention,
            'output_layer':self.output_layer
            })
        return config

Code for compilation and fitting :

def adapts_compile_fit(model,
                       train_data,
                       valid_data,
                       n_epochs = 40,
                       n_patience = 20,
                       init_lr = 0.001,
                       lr_decay_rate = 0.1):

  # early_stopping = keras.callbacks.EarlyStopping(monitor="val_accuracy",patience = n_patience,restore_best_weights = True)
  n_decay_steps = n_epochs * len(list(train_ds))
  scheduled_lr = keras.optimizers.schedules.ExponentialDecay(
        initial_learning_rate=init_lr,
        decay_steps=n_decay_steps,
        decay_rate=lr_decay_rate)
  model.compile(loss='sparse_categorical_crossentropy',
                optimizer = keras.optimizers.RMSprop(learning_rate = scheduled_lr),
                metrics = ['accuracy'])
  return model.fit(train_ds,
                   epochs = n_epochs,
                   validation_data = valid_ds
                   ,verbose=1)

I have saved it using model.save("/content/tf.keras") and loading the model as my_model_2 = keras.models.load_model("/content/tf.keras",custom_objects={'BidirectionalEncoderandDecoderWithAttention':BidirectionalEncoderandDecoderWithAttention}) and getting following error :

TypeError: Unable to revive model from config. When overriding the get_config() method, make sure that the returned config contains all items used as arguments in the constructor to <class ‘main.BidirectionalEncoderandDecoderWithAttention’>, which is the default behavior. You can override this default behavior by defining a from_config(cls, config) class method to specify how to create an instance of BidirectionalEncoderandDecoderWithAttention from its config.

Received config={‘input_vocabulary_size’: 14617, ‘output_vocabulary_size’: 29604, ’embedding_size’: 256, ‘units_lstm’: 512, ‘encoder_embedding’: {‘module’: ‘keras.layers’, ‘class_name’: ‘Embedding’, ‘config’: {‘name’: ’embedding_6′, ‘trainable’: True, ‘dtype’: ‘float32’, ‘batch_input_shape’: [None, None], ‘input_dim’: 14617, ‘output_dim’: 256, ’embeddings_initializer’: {‘module’: ‘keras.initializers’, ‘class_name’: ‘RandomUniform’, ‘config’: {‘minval’: -0.05, ‘maxval’: 0.05, ‘seed’: None}, ‘registered_name’: None}, ’embeddings_regularizer’: None, ‘activity_regularizer’: None, ’embeddings_constraint’: None, ‘mask_zero’: True, ‘input_length’: None}, ‘registered_name’: None, ‘build_config’: {‘input_shape’: [None, None]}}, ‘encoder’: {‘module’: ‘keras.layers’, ‘class_name’: ‘Bidirectional’, ‘config’: {‘name’: ‘bidirectional_3’, ‘trainable’: True, ‘dtype’: ‘float32’, ‘layer’: {‘module’: ‘keras.layers’, ‘class_name’: ‘LSTM’, ‘config’: {‘name’: ‘lstm_6’, ‘trainable’: True, ‘dtype’: ‘float32’, ‘return_sequences’: True, ‘return_state’: True, ‘go_backwards’: False, ‘stateful’: False, ‘unroll’: False, ‘time_major’: False, ‘units’: 256, ‘activation’: ‘tanh’, ‘recurrent_activation’: ‘sigmoid’, ‘use_bias’: True, ‘kernel_initializer’: {‘module’: ‘keras.initializers’, ‘class_name’: ‘GlorotUniform’, ‘config’: {‘seed’: None}, ‘registered_name’: None, ‘shared_object_id’: 138063236840848}, ‘recurrent_initializer’: {‘module’: ‘keras.initializers’, ‘class_name’: ‘Orthogonal’, ‘config’: {‘gain’:…

Error encountered during deserialization: (‘Keyword argument not understood:’, ‘encoder_embedding’)

I have added all the parameters in config file but still error has not gone. Please Help??