I am loading the network like this:
def load_network(model_path, device): 
    if 'cuda' in device:
        torch.backends.cudnn.benchmark = True
    else:
        torch.backends.cudnn.benchmark = False
            
    inpainting_model = PConvUNet().to(device)
    inpainting_model.eval()
    load_ckpt(model_path, [('model', inpainting_model)], device=device)

def load_ckpt(ckpt_name, models, optimizers=None, device=None):
    ckpt_dict = torch.load(ckpt_name, map_location=device)
    for prefix, model in models:
        assert isinstance(model, nn.Module)
        model.load_state_dict(ckpt_dict[prefix], strict=False)
    if optimizers is not None:
        for prefix, optimizer in optimizers:
            optimizer.load_state_dict(ckpt_dict[prefix])
    return ckpt_dict['n_iter']

and then just calling jit.trace:
net = load_network(net_path, device)
#dummy input for tracing the model
image = torch.rand(1, 3, inpainting_size[0], inpainting_size[1])
mask = torch.rand(1, 3, inpainting_size[0], inpainting_size[1])
                
image = image.to(device)
mask = mask.to(device)
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(net, (image, mask))
traced_script_module.save("model/InpaintNet_cpp.pt")