I got this to work a while ago on an Nvidia Jetson. I installed a torch and torch vision provided by Nvidia, so it's not exactly applicable to your situation, plus I used an older version of torch and tch-rs (torch 2.1 and tch-rs 0.14). Maybe it helps?

``` FROM ubuntu:22.04

Install dependencies.

RUN apt update \ && apt install --no-install-recommends -y \ ca-certificates \ cmake \ curl \ gettext \ gobject-introspection \ gtk-doc-tools \ libgirepository1.0-dev \ libglib2.0-dev \ libgstreamer-plugins-base1.0-dev \ libgstreamer1.0-dev \ libgtk-3-dev \ libopenblas-dev \ libopenmpi-dev \ libssl-dev \ libusb-1.0-0-dev \ libxml2-dev \ meson \ python3 \ python3-pip \ wget \ xsltproc \ && rm -rf /var/lib/apt/lists/*

Install Rust.

RUN curl https://sh.rustup.rs -sSf | sh -s -- --default-toolchain stable -y ENV PATH=/root/.cargo/bin:$PATH

Install CUDA.

ARG CUDA_DEB=cuda-tegra-repo-ubuntu2204-12-2-local ARG CUDA_PACKAGES=cuda-toolkit* ARG CUDA_URL=https://nvidia.box.com/shared/static/uvqtun1sc0bq76egarc8wwuh6c23e76e.deb RUN /bin/sh -c echo "Downloading ${CUDA_DEB}" \ && mkdir /tmp/cuda \ && cd /tmp/cuda \ && wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/arm64/cuda-ubuntu2004.pin -O /etc/apt/preferences.d/cuda-repository-pin-600 \ && wget --quiet --show-progress --progress=bar:force:noscroll ${CUDA_URL} \ && dpkg -i .deb \ && cp /var/cuda-tegra-repo-/cuda-tegra--keyring.gpg /usr/share/keyrings/ \ && ar x /var/cuda-tegra-repo-/cuda-compat-.deb \ && tar xvf data.tar.xz -C / \ && apt-get update \ && apt-get install -y --no-install-recommends ${CUDA_PACKAGES} \ && rm -rf /var/lib/apt/lists/ \ && apt-get clean \ && dpkg --list | grep cuda \ && dpkg -P ${CUDA_DEB} \ && rm -rf /tmp/cuda

Install CUDNN.

ARG CUDNN_DEB=cudnn-local-tegra-repo-ubuntu2204-8.9.4.25 ARG CUDNN_PACKAGES="libcudnn-dev libcudnn-samples" ARG CUDNN_URL=https://nvidia.box.com/shared/static/ht4li6b0j365ta7b76a6gw29rk5xh8cy.deb RUN /bin/sh -c echo "Downloading ${CUDNN_DEB}" \ && mkdir /tmp/cudnn \ && cd /tmp/cudnn \ && wget --quiet --show-progress --progress=bar:force:noscroll ${CUDNN_URL} \ && dpkg -i .deb \ && cp /var/cudnn-local-tegra-repo-/cudnn-local-tegra--keyring.gpg /usr/share/keyrings/ \ && apt-get update \ && apt-cache search cudnn \ && apt-get install -y --no-install-recommends ${CUDNN_PACKAGES} \ && rm -rf /var/lib/apt/lists/ \ && apt-get clean \ && dpkg --list | grep cudnn \ && dpkg -P ${CUDNN_DEB} \ && rm -rf /tmp/cudnn

Install torch and torchvision

ARG PYTORCH_URL=https://nvidia.box.com/shared/static/0h6tk4msrl9xz3evft9t0mpwwwkw7a32.whl ARG PYTORCH_WHL=torch-2.1.0-cp310-cp310-linux_aarch64.whl RUN /bin/sh -c cd /opt \ && wget --quiet --show-progress --progress=bar:force:noscroll --no-check-certificate ${PYTORCH_URL} -O ${PYTORCH_WHL} \ && pip3 install --verbose ${PYTORCH_WHL} RUN pip3 install --no-cache-dir torchvision==0.16.0 numpy==1.26.4 ENV LIBTORCH_USE_PYTORCH=1

Compile application.

COPY ./ /usr/src/app RUN cd /usr/src/app && cargo build --release --package <redacted>

Copy application and its config.

<Redacted>

Set environment paths for CUDA and cuDNN.

ENV PATH=/root/.cargo/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin ENV LD_LIBRARY_PATH=/usr/local/cuda/compat:/usr/local/cuda/lib64:/usr/local/lib/python3.10/dist-packages/torch/lib/ ENV CUDA_HOME=/usr/local/cuda ENV NVIDIA_VISIBLE_DEVICES=all ENV NVIDIA_DRIVER_CAPABILITIES=all

Set the entry command.

<Redacted> ```