Hi Sebastian. I am glad to know the code is being useful. To save the checkpoint is really straightforward. Note that UNET inherits from nn.Module, so it behaves like a Pytorch module and it is the actual model. To save the weights you will use torch.save(model.state_dict(), filename). In the example, the model is called unet (lowercase). All the information can be found here: https://pytorch.org/tutorials/beginner/saving_loading_models.html

Regards,

Mauricio