https://docs.pytorch.org/docs/stable/export.html
import torch
import torchvision.models as models
net = models.resnet18()
net.eval()
torch.manual_seed(0)
x = torch.rand(1, 3, 224, 224)
a = net(x)
# export pt2
program = torch.export.export(net, (x,))
torch.export.save(program, 'test_resnet18.pt2')