CUDA not detected when building Torch Examples on DeepStream Docker container

Exact reproduction steps:

  1. Run the docker container: sudo docker run -it --rm --net=host --runtime nvidia -e DISPLAY=$DISPLAY -w /opt/nvidia/deepstream/deepstream-5.0 -v /tmp/.X11-unix/:/tmp/.X11-unix nvcr.io/nvidia/deepstream-l4t:5.0-dp-20.04-samples

  2. Install git and cmake: apt-get update && apt-get install git cmake

  3. Export CUDA Home variable: export CUDA_HOME=/usr/local/cuda

  4. Install Torch from ARM Wheel that’s provided by NVIDIA as documented here

wget https://nvidia.box.com/shared/static/3ibazbiwtkl181n95n9em3wtrca7tdzp.whl -O torch-1.5.0-cp36-cp36m-linux_aarch64.whl
sudo apt-get update
sudo apt-get -qq install libopenblas-base libopenmpi-dev python3-pip
pip3 install Cython
pip3 install numpy torch-1.5.0-cp36-cp36m-linux_aarch64.whl
  1. Clone the examples repository: cd ~ && mkdir Github && git clone https://github.com/pytorch/examples

  2. Run cmake: cd examples/cpp/mnist && mkdir build && cd build && cmake -DCMAKE_PREFIX_PATH=/usr/local/lib/python3.6/dist-packages/torch ..

Output:

 cmake -DCMAKE_PREFIX_PATH=/root/.local/lib/python3.6/site-packages/torch/ ..
-- Looking for pthread.h
-- Looking for pthread.h - found
-- Looking for pthread_create
-- Looking for pthread_create - not found
-- Looking for pthread_create in pthreads
-- Looking for pthread_create in pthreads - not found
-- Looking for pthread_create in pthread
-- Looking for pthread_create in pthread - found
-- Found Threads: TRUE  
-- Could NOT find CUDA (missing: CUDA_CUDART_LIBRARY) (found version "10.0")
CMake Warning at /root/.local/lib/python3.6/site-packages/torch/share/cmake/Caffe2/public/cuda.cmake:31 (message):
  Caffe2: CUDA cannot be found.  Depending on whether you are building Caffe2
  or a Caffe2 dependent library, the next warning / error will give you more
  info.
Call Stack (most recent call first):
  /root/.local/lib/python3.6/site-packages/torch/share/cmake/Caffe2/Caffe2Config.cmake:88 (include)
  /root/.local/lib/python3.6/site-packages/torch/share/cmake/Torch/TorchConfig.cmake:40 (find_package)
  CMakeLists.txt:5 (find_package)


CMake Error at /root/.local/lib/python3.6/site-packages/torch/share/cmake/Caffe2/Caffe2Config.cmake:90 (message):
  Your installed Caffe2 version uses CUDA but I cannot find the CUDA
  libraries.  Please set the proper CUDA prefixes and / or install CUDA.
Call Stack (most recent call first):
  /root/.local/lib/python3.6/site-packages/torch/share/cmake/Torch/TorchConfig.cmake:40 (find_package)
  CMakeLists.txt:5 (find_package)


-- Configuring incomplete, errors occurred!
See also "/root/Github/examples/cpp/mnist/build/CMakeFiles/CMakeOutput.log".
See also "/root/Github/examples/cpp/mnist/build/CMakeFiles/CMakeError.log".

This issue was also a problem for me with the pytorch container.

Have you added “export PATH=/usr/local/cuda/bin:$PATH” and “export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH” to the .bashrc file?

Yes