pytorch-cuda.sh

#!/bin/sh

set -x

rm -rf /etc/apt/sources.list.d/nvidia-l4t-apt-source.list
echo "deb https://repo.download.nvidia.com/jetson/common r32.4 main" > /etc/apt/sources.list.d/nvidia-l4t-apt-source.list
echo "deb https://repo.download.nvidia.com/jetson/t194 r32.4 main" >> /etc/apt/sources.list.d/nvidia-l4t-apt-source.list
apt-key adv --fetch-key https://repo.download.nvidia.com/jetson/jetson-ota-public.asc
apt update

git clone --recursive --branch v1.7.0 http://github.com/pytorch/pytorch
cd pytorch
wget https://gist.githubusercontent.com/dusty-nv/ce51796085178e1f38e3c6a1663a93a1/raw/9d7261584a7482e7cc0fcb08a4a232c6d023f812/pytorch-1.7-jetpack-4.4.1.patch
git apply pytorch-1.7-jetpack-4.4.1.patch


export USE_NCCL=0
export USE_DISTRIBUTED=0
export USE_QNNPACK=0
export USE_PYTORCH_QNNPACK=0
export TORCH_CUDA_ARCH_LIST="7.2"
export PYTORCH_BUILD_VERSION=1.7.0
export PYTORCH_BUILD_NUMBER=1

# Install CUDA libs
sudo apt install libcudnn8-dev cuda-toolkit-10-2 -y

sudo apt-get install python3-pip cmake libopenblas-dev python3-venv -y

python3 -m venv venv
source venv/bin/activate
pip install wheel cython pyyaml
pip3 install -r requirements.txt
pip3 install scikit-build
pip3 install ninja
pip3 install dataclasses
python3 setup.py bdist_wheel