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