PyTorch can't find CUDA inside JetPack 5.1 docker container

[divelix]

Platform: Xavier NX with JetPack 4.4 [L4T 32.4.2] installed
I pulled latest (5.1) Docker image from here:

NVIDIA L4T PyTorch | NVIDIA NGC

Inside container (which I run by official instruction in the link above, using --runtime=nvidia), I get the following output from python3 commands:
torch.__version__ outputs

2.0.0a0+ec3941ad.nv23.02

torch.cuda.is_available() outputs

/usr/local/lib/python3.8/dist-packages/torch/cuda/__init__.py:98: UserWarning: CUDA initialization: CUDA driver initialization failed, you might not have a CUDA gpu. (Triggered internally at /home/riship/old_pyt/pytorch/c10/cuda/CUDAFunctions.cpp:109.)
  return torch._C._cuda_getDeviceCount() > 0
False

Meanwhile, CUDA seems to be installed, as ls /usr/local outputs

bin  cuda  cuda-10.2  cuda-11  cuda-11.4  etc  games  include  lib  man  sbin  share  src

Why PyTorch can’t see CUDA? Can I fix that somehow?

[divelix]

I tried other image versions and it seems, like host’s and container JetPack versions should coinside exactly. On my platform, 32.4.2-pth1.5 image outputs True for torch.cuda.is_available().

[dusty_nv]

@divelix

On my platform, 32.4.2-pth1.5 image outputs True for torch.cuda.is_available().

Hi @divelix, yes you are correct - JetPack4 won’t run containers built for JetPack5 (and vice versa). Once on JetPack5, those containers can run on other minor versions of JetPack. But since JetPack4 mounts the CUDA/cuDNN libraries from the host, those containers should match the version of L4T you’re running if you’re on JetPack4.