How to save the checkpoint for self-defined model?

Hi,
I am interested about saving the checkpoint of my self-defined model by subclassing `snt.Module`, but I couldn't find solid resource online. Can I use default method in tensorflow 2 for keras model like `tf.keras.callbacks.ModelCheckpoint()`? Does this work for sonnet 2 model?

Thanks,
ZY

Start by checking Sonnet 2's guidance for saving checkpoints from a self-defined snt.Module and comparing it with TensorFlow 2's tf.keras.callbacks.ModelCheckpoint. Done means providing a clear, tested explanation of whether the Keras callback works and what checkpointing approach Sonnet users should follow.

Written by the indexing model from the issue text.