Save model as h5 in pix2pix tutorial example

Hi,

I am following the tutorial on pix2pix pix2pix: tradução de imagem para imagem com um GAN condicional  |  TensorFlow Core
and have trouble generating new images after trainingfrom checkpoint restore…so I thought that saving a model as h5 could be a solution. I found the method tf.keras.callbacks.ModelCheckpoint, however, I read, it can only work in conjunction with model.fit(). and this example utilizes custom solutions. Can anyone advise how to combine those custom and prebuilt functions together? Any useful tutorials or examples?

Best,
Maria