Install PyTorch for CUDA 9.0

env.sh

PREFIX=$PWD/env

export CUDA_HOME=/usr/local/cuda-9.0
export CMAKE_PREFIX_PATH=$PREFIX
export LD_LIBRARY_PATH=$PREFIX/lib:$LD_LIBRARY_PATH
export CPATH=$PREFIX/include:$CPATH
export PATH=$PREFIX/bin:$CUDA_HOME/bin:${PATH//"/home/cvlab/anaconda3/bin"/}
export CUDNN_LIB_DIR=$PREFIX/lib
export CUDNN_INCLUDE_DIR=$PREFIX/include

if [ "$1" = "install" ]; then
        export TMPDIR=$PREFIX/tmpdir
        export MAKEFLAGS=-j16
        mkdir -p "$TMPDIR"
        set -ex
        wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh
        bash Miniconda3-latest-Linux-x86_64.sh -b -u -p "$PREFIX"
        rm Miniconda3-latest-Linux-x86_64.sh
        conda install -y numpy pyyaml mkl mkl-include setuptools cmake cffi typing mkl-dnn opencv cudnn=7.3.1=cuda9.0_0
        conda install -y -c pytorch magma-cuda90
        git clone --recursive https://github.com/pytorch/pytorch "$PREFIX/pytorch"
        cd "$PREFIX/pytorch"
        TORCH_CUDA_ARCH_LIST=6.1 python3 setup.py install
        pip install torchvision IPython h5py seaborn
        rm -rf "$TMPDIR"
elif [ "$1" = "in" ]; then
        bash --rcfile <(cat "$HOME/.bashrc"; cat env.sh; echo 'export PS1="$PS1 ENV $ "') -i
fi