No CUDA runtime is found, using CUDA_HOME=‘/usr/local/cuda’

Hi

Previously, I create Python virtual environment in my Jetson Orin Nano (my CUDA version is 11.4.315) by using torch = 2.1.0a0+41361538.nv23.06 and torchvision = 0.16.0+fbb4cc5. It can work well without any error both installation process and inference. After that try to do the same thing with docker, I pulled the l4t-tensorrt image:nvcr.io/nvidia/l4t-tensorrt:r8.5.2.2-devel. My docker image is built successfully and torch.cuda.is_available() returned True after I installed torch, but when I installed torchvision I got this error: No CUDA runtime is found, using CUDA_HOME=‘/usr/local/cuda’

.

This is my code inside bash script to install torchvision:

 sudo apt install -y libjpeg-dev zlib1g-dev
 git clone https://github.com/pytorch/vision torchvision
 cd /home/torch/torchvision
 git checkout v0.16.0
 python3 setup.py install --user

This is my DockerFile:

# Use NVIDIA L4T TensorRT image as the base
FROM nvcr.io/nvidia/l4t-tensorrt:r8.5.2.2-devel

# Avoid prompts from apt
ENV DEBIAN_FRONTEND=noninteractive

# Set CUDA paths
ENV CUDA_HOME=/usr/local/cuda
ENV PATH=/usr/local/cuda/bin:${PATH}
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH}

# Set the working directory in the container
WORKDIR /usr/src/app

# Copy the Python script into the container
COPY setup_docker.sh /usr/src/app/setup_docker.sh

RUN chmod +x /usr/src/app/setup_docker.sh && /bin/bash /usr/src/app/setup_docker.sh

I guess my problem is CUDA paths setting …
How can I set the correct CUDA paths or do you have a complete DockerFile to setup and install torchvision with nvcr.io/nvidia/l4t-tensorrt:r8.5.2.2-devel ?

Thnaks