Subclassed model not saved by ModelCheckpoint callback

Hi,

I’m using Keras and TensorFlow 2.12.0 to implement a subclassed model with custom methods, such as a custom train_step. To save the model, I use the ModelCheckpoint callback.

My goal is to save not only the model’s architecture and weights but also the custom methods and training configuration, allowing me to resume training from the exact state where it was left off. From the Keras/TensorFlow documentation, it seems that the HDF5 format can only save the weights of a subclassed model. Therefore, I tried to use the SavedModel and Keras formats. However, after loading the model, the training configuration and custom methods are not available.

Is there something that I might be missing?
Working code snippet to reproduce the issue: Google Colab