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



Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store