saving checkpoints in "model.pth" format and visualizing the segmentation results

[k-nayak]

Hello @LeeJunHyun , thank you for the amazing code. I want to use Attention Unet for real-time segmentation and see how well it works. I need weight files for using the model with a camera setup usually in ".pth" format so i can feed new images and see the segmentation results. Is there any way to save the checkpoint or is the ".pkl" something that can be used instead. Thanks in advance.

[Lloyd-Pottiger]

The Pytorch offical tutorials may can help you https://pytorch.org/tutorials/beginner/basics/saveloadrun_tutorial.html, and you can just change the unet_path with '.pth'.