Converting UNET to Torchscript format for AWS Lambda


According to AWS Lambda deployment instructions on FASTAI, the SAM application expects a PyTorch model in TorchScript format to be saved to S3 along with a classes text file with the output class names.

Previously, I tried to use Torchscript with UNET and got errors due to its hooks. Has someone tried this deployment in the context of UNET or just UNET to torchscript conversion?