AdamLRM fails to be saved/loaded inside a model

[llouislu]

The model compiled with AdamLRM fails to be saved, complaining the following error:

AttributeError: 'AdamLRM' object has no attribute 'lr_multiplier'

I have looked into the code, the AdamLRM class doesn't hold the lr_multiplier dict for get_config method.
As a quick fix, add self.lr_multiplier = lr_multiplier in __init__.

However, the model fails to load after the fix, prompting ValueError: Unknown optimizer: AdamLRM.

Code to reproduce

import tensorflow as tf
tf.config.set_visible_devices([], "GPU")

from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from adamlrm import AdamLRM
import numpy as np
X = np.zeros([5,100])
Y = np.zeros([5])

model = Sequential()
model.add(Dense(64, activation='relu', input_dim=100))
model.add(Dense(1, activation='softmax'))

lr_multiplier = {
  'var1':1e-2, # optimize 'var1*' with a smaller learning rate
  'var2':10   # optimize 'var2*' with a larger learning rate
  }
  
opt = AdamLRM(lr=0.001, lr_multiplier=lr_multiplier)

model.compile(
  optimizer=opt,
  loss='mse',
  metrics=['mse'])

model.save("test_model")


del model

model = tf.keras.models.load_model("test_model")
print(model.summary())

[akionux]

Corrected registering multipliers in get_config method.

After the commit, saving model works, but loading the model still fails.
According to Keras FAQ,
model = tf.keras.models.load_model("test_model", custom_objects={'AdamLRM':AdamLRM})
should work, but does not.

By loading uncompiled model and compiling it again, it works.

model = tf.keras.models.load_model("test_model", compile=False)
model.compile(
  optimizer=opt,
  loss='mse',
  metrics=['mse'])