How to get cuda working in a docker container for pytorch applications

Autonomous Machines Jetson & Embedded Systems Jetson AGX Orin
,
1

I’m trying to run a pytorch application in docker on a Jetson AGX Orin but I’m having trouble getting it to work with cuda (see code-block 1). Torch works fine when running on the host (code-block 2). I’ve looked around online for quite a while but I’m stuck.

Which container image should I use? And is there any more configuration that needs to be done to get it to work?

$ docker run --runtime nvidia --rm -it --network=host nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3 bash

root@jetson:/# python3
Python 3.8.10 (default, Nov 14 2022, 12:59:47) 
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
False

$ python3
Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True

Jetson info:

$ sudo apt-cache show nvidia-jetpack
[sudo] password for move: 
Package: nvidia-jetpack
Source: nvidia-jetpack (6.0)
Version: 6.0+b106
Architecture: arm64
Maintainer: NVIDIA Corporation
Installed-Size: 194
Depends: nvidia-jetpack-runtime (= 6.0+b106), nvidia-jetpack-dev (= 6.0+b106)
Homepage: http://developer.nvidia.com/jetson
Priority: standard
Section: metapackages
Filename: pool/main/n/nvidia-jetpack/nvidia-jetpack_6.0+b106_arm64.deb
Size: 29296
SHA256: 561d38f76683ff865e57b2af41e303be7e590926251890550d2652bdc51401f8
SHA1: ef3fca0c1b5c780b2bad1bafae6437753bd0a93f
MD5sum: 95de21b4fce939dee11c6df1f2db0fa5
Description: NVIDIA Jetpack Meta Package
Description-md5: ad1462289bdbc54909ae109d1d32c0a8

Package: nvidia-jetpack
Source: nvidia-jetpack (6.0)
Version: 6.0+b87
Architecture: arm64
Maintainer: NVIDIA Corporation
Installed-Size: 194
Depends: nvidia-jetpack-runtime (= 6.0+b87), nvidia-jetpack-dev (= 6.0+b87)
Homepage: http://developer.nvidia.com/jetson
Priority: standard
Section: metapackages
Filename: pool/main/n/nvidia-jetpack/nvidia-jetpack_6.0+b87_arm64.deb
Size: 29298
SHA256: 70be95162aad864ee0b0cd24ac8e4fa4f131aa97b32ffa2de551f1f8f56bc14e
SHA1: 36926a991855b9feeb12072694005c3e7e7b3836
MD5sum: 050cb1fd604a16200d26841f8a59a038
Description: NVIDIA Jetpack Meta Package
Description-md5: ad1462289bdbc54909ae109d1d32c0a8
3

I got it working with this container dustynv/l4t-pytorch:r36.2.0

1 Like
4

OK yes @helgi.bjornsson, glad you got it working - those images from jetson-containers are updated for JetPack 6.

1 Like
5

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