Pytorch : Unable to Train YoloV5 on Jetson Orin AGX GPU (using CUDA)

Device : Jetson Orin AGX
Jetpack Version : 5.1
CUDA Version : 11.4

So it went well when I trained my Obj detection with Yolo by using my CPU (with torch 2.0.1 from pytorch website).

However when I tried to train my Obj detection by using my GPU, it failed after just one epoch. Heres the Runtime error:

RuntimeError: Couldn't load custom C++ ops. This can happen if your PyTorch and torchvision versions are incompatible, 
or if you had errors while compiling torchvision from source. For further information on the compatible versions, check for the compatibility matrix. 
Please check your PyTorch version with torch.__version__ and your torchvision version with torchvision.__version__ and verify if they are compatible, and if not please reinstall torchvision so that it matches your PyTorch install.

Notice that it might caused because my torch and torchvision versions are incompatible.

I followed Pytorch Installation on Jetson platfrom from Installing PyTorch for Jetson Platform Documentation. There was no error during the installation of both libraries.

Here is the log when I printed torch and torchvision versions:

@ubuntu:~$ python3
Python 3.8.10 (default, Mar 13 2023, 10:26:41)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.__version__
>>> torch.cuda.is_available()
>>> torch.cuda.current_device()
>>> torch.cuda.get_device_name(0)
>>> import torchvision
/home/happibot/.local/lib/python3.8/site-packages/torchvision/io/ UserWarning: Failed to load image Python extension: ''If you don't plan on using image functionality from ``, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?
>>> torchvision.__version__ 

Both versions are compatible according to GitHub - pytorch/vision: Datasets, Transforms and Models specific to Computer Vision.

However it is so unfortunate that there is no guide to install torchvision for pytorch version built by NVIDIA. Why does this happen?


Could you try TorchVision v0.14.1 instead?
Or you can try below a container that has PyTorch and TorchVision installed.


This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.