I subclassed a Keras layer from tf.keras.layers.AbstractRNNCell and created a custom RNN layer with a set of weights and states. I then passed this into a custom Keras model class (subclassed from tf.keras.Models). However, after training the model, when I try to go model.save('my_model', save_format='tf' ) and `tf.keras.models.load_model(‘my_model’)', I get the following error
ValueError: Unable to restore custom object of type _tf_keras_rnn_layer. Please make sure that any custom layers are included in the custom_objects arg when calling load_model() and make sure that all layers implement get_config and from_config.