Pytorch with meta-tegra jetpack6

[dlheard17]

Wondering if there is a proven solution for using Pytorch with a meta-tegra based yocto image on jetpack6?
I'm using scarthgap(python3.12) and noticed that Jetpack6.2 does not yet have a distribution for pytorch here https://developer.download.nvidia.com/compute/redist/jp/

I've tried using the non-jetson-specific versions of torch and cuda is broken with that config.

python3 snippet:

import torch
torch.cuda.is_available()
True
torch.zeros(5, dtype=torch.complex64, device="cuda")
Traceback (most recent call last):
File "", line 1, in
RuntimeError: CUDA error: no kernel image is available for execution on the device
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.

torch.zeros(5, dtype=torch.complex64)
tensor([0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j])

[madisongh]

I'm not aware of a CUDA-enabled recipe for pytorch that you can use with OE/Yocto builds.

Searching around, I see that there may be an NVIDIA pre-built copy of pytorch compatible with JetPack 6.2 - dev forum post 1 and dev forum post 2 - but I'm not sure it would be usable in a Yocto-built OS, due to python differences.

You might look into using the NVIDIA docker containers instead of trying to integrated pytorch directly in your OS image. I believe that's a commonly used approach to solving this problem.

[sefabre]

Hello,
I am working on orin nano nx module in P3768 carrier with scarthgap branch (using meta-tegrademo). So, CUDA version is 12.6 and python3 version is 3.12
I would like to integrate pytorch, torchvision and onnxruntime python3 modules but I do not find corresponding recipes:

• on meta-python-ai, they do not use cuda
• on master branch of meta-tegra-community, recipes do not generate python3 modules

Moreover, the prebuild modules in https://pypi.jetson-ai-lab.io/jp6/cu126 or in https://github.com/ultralytics/assets/releases are generated for python 3.10 (and not 3.12)

I have tried to retrieve modules from NVIDIA pytorch docker: https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-12.html (the last version with CUDA 12.6 and python3.12). But torchvision version (0.20.0) is not compatible with the torch (2.6.0) (I do not know if it is why I have the following error:"RuntimeError: CUDA error: no kernel image is available for execution on the device"). Same error inside the docker (or with release 24.11 for docker)

[madisongh]

I think it should be pretty straightforward to have the pytorch recipe we have in meta-tegra-community to build the python bindings, too. @ichergui what do you think?

[ichergui]

Hi @sefabre
We have Pytorch, torchvision and onnxruntime recipes in meta-tegra-community (master branch).

• onnxruntime recipe provides both C/C++ and python bindings libraries.
• Pytorch and Torchvision only C/C++ libraries are available.

I agree @madisongh. I will try to enable the python bindings for Pytorch and Torchvivion.
@sefabre contributions are more than welcome.

[ichergui]

Hey @dlheard17 @sefabre

Pytorch with Python bindings is available now in our branches master and master-l4t-r38.4.x.

I'm planing to enable the Python bindings in onnx

Please try it and let me know how it goes.

Thanks to @pberginkonsult for his contributions.