Can't save custom subclassed model

Question

Inspired by tf.keras.Model subclassing I created custom model.
I can train it and get successfull results, but I can\'t save it.
I use python3.6 wit

Answer 1

Actually recreating the model with

keras.models.load_model('path_to_my_model')

didn't work for me

First we have to save_weights from the built model

model.save_weights('model_weights', save_format='tf')

Then we have to initiate a new instance for the subclass Model then compile and train_on_batch with one record and load_weights of built model

loaded_model = ThreeLayerMLP(name='3_layer_mlp')
loaded_model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
loaded_model.train_on_batch(x_train[:1], y_train[:1])
loaded_model.load_weights('model_weights')

This work perfectly in TensorFlow==2.2.0