PyTorch & Torchvision compatible issue on Jetson Orin

🐛 Describe the bug

Previous discussion: https://forums.developer.nvidia.com/t/pytorch-torchversion-compatible-issue-on-l4t35-5-0/310929/9

daniel@daniel-nvidia:~/Work/yolov5$ python detect.py --weights yolov5s.pt --source ../../Videos/Worlds_longest_drone_fpv_one_shot.mp4
WARNING ⚠️ Python>=3.10 is required, but Python==3.8.10 is currently installed
/home/daniel/.local/lib/python3.8/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: '/home/daniel/.local/lib/python3.8/site-packages/torchvision/image.so: undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKSsb'If you don't plan on using image functionality from `torchvision.io`, 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?
  warn(
detect: weights=['yolov5s.pt'], source=../../Videos/Worlds_longest_drone_fpv_one_shot.mp4, data=data/coco128.yaml, imgsz=[640, 640], conf_thres=0.25, iou_thres=0.45, max_det=1000, device=, view_img=False, save_txt=False, save_format=0, save_csv=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs/detect, name=exp, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False, dnn=False, vid_stride=1
YOLOv5 🚀 v7.0-378-g2f74455a Python-3.8.10 torch-2.1.0a0+41361538.nv23.06 CUDA:0 (Orin, 7451MiB)

Fusing layers...
YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients
Traceback (most recent call last):
  File "detect.py", line 437, in <module>
    main(opt)
  File "detect.py", line 432, in main
    run(**vars(opt))
  File "/home/daniel/.local/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "detect.py", line 210, in run
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
  File "/home/daniel/Work/yolov5/utils/general.py", line 1104, in non_max_suppression
    i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
  File "/home/daniel/.local/lib/python3.8/site-packages/torchvision/ops/boxes.py", line 40, in nms
    _assert_has_ops()
  File "/home/daniel/.local/lib/python3.8/site-packages/torchvision/extension.py", line 46, in _assert_has_ops
    raise RuntimeError(
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 https://github.com/pytorch/vision#installation 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.

Versions

daniel@daniel-nvidia:~/Work/yolov5$ python -c "import torch; import torchvision; print(f'PyTorch version: {torch.__version__}'); print(f'Torchvision version: {torchvision.__version__}')"
/home/daniel/.local/lib/python3.8/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: '/home/daniel/.local/lib/python3.8/site-packages/torchvision/image.so: undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKSsb'If you don't plan on using image functionality from `torchvision.io`, 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?
  warn(
PyTorch version: 2.1.0a0+41361538.nv23.06
Torchvision version: 0.16.1+fdea156

daniel@daniel-nvidia:~/Work$ python collect_env.py
Collecting environment information...
PyTorch version: 2.1.0a0+41361538.nv23.06
Is debug build: False
CUDA used to build PyTorch: 11.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (aarch64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Sep 11 2024, 16:02:53)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.10.192-tegra-aarch64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.4.315
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/aarch64-linux-gnu/libcudnn.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_adv_infer.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_adv_train.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_cnn_infer.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_cnn_train.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_ops_infer.so.8.6.0
/usr/lib/aarch64-linux-gnu/libcudnn_ops_train.so.8.6.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False

CPU:
Architecture:                       aarch64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
CPU(s):                             6
On-line CPU(s) list:                0-5
Thread(s) per core:                 1
Core(s) per socket:                 3
Socket(s):                          2
Vendor ID:                          ARM
Model:                              1
Model name:                         ARMv8 Processor rev 1 (v8l)
Stepping:                           r0p1
CPU max MHz:                        1510.4000
CPU min MHz:                        115.2000
BogoMIPS:                           62.50
L1d cache:                          384 KiB
L1i cache:                          384 KiB
L2 cache:                           1.5 MiB
L3 cache:                           2 MiB
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; CSV2, but not BHB
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected
Flags:                              fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm lrcpc dcpop asimddp uscat ilrcpc flagm

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] onnx==1.17.0
[pip3] onnx-graphsurgeon==0.3.12
[pip3] onnxruntime-gpu==1.17.0
[pip3] onnxslim==0.1.35
[pip3] torch==2.1.0a0+41361538.nv23.6
[pip3] torchvision==0.16.1+fdea156
[conda] Could not collect

Start with detect.py and utils/general.py around the torchvision.ops.nms call, then compare the reported PyTorch 2.1.0a0 and torchvision 0.16.1 versions with the compatibility matrix linked in the error. Reproduce the command on the reported Jetson Orin environment; done means detection runs without the custom C++ ops error.