Converting a .pt file to a .engine file on Jetson Xavier NX


I have my own pretrained pytorch model that I want to convert to a TensorRT model (.engine), I run this python script:

import torch
from torch2trt import torch2trt

pytorch_model = torch.load(“”)
x = torch.ones((1, 3, 224, 224)).cuda()
trt_model = torch2trt(pytorch_model, [x])
with open(“model.engine”, “wb”) as f:

but everytime I run it, the following error occurs:

Traceback (most recent call last):
File “/home/nvidia/torch2trt/”, line 9, in pytorch_model = torch.load(“”)*
*File “/home/nvidia/.local/lib/python3.6/site-packages/torch/”, line 592, in load return_load(opened_zipfile, map_location, pickle_module, *pickle_load_args)
File “/home/nvidia/.local/lib/python3.6/site-packages/torch/”, line 851, in _load
result = unpickler.load()
ModuleNotFoundError: No module named ‘models’

How can I solve it?


TensorRT Version:
GPU Type: Xavier
CUDA Version: 10.2
Operating System + Version: Ubuntu 18.04
Python Version (if applicable): 3.6.9
PyTorch Version (if applicable): 1.8.0


This looks like a Jetson issue. Please refer to the below samples in case useful.

For any further assistance, we will move this post to to Jetson related forum.