Install PyTorch with currently available system CUDA version

get_torch_with_cuda.sh

#get cuda version and install torch

if nvcc --version 2&> /dev/null; then
    # Determine CUDA version using default nvcc binary
    CUDA_VERSION=$(nvcc --version | sed -n 's/^.*release \([0-9]\+\.[0-9]\+\).*$/\1/p');

elif /usr/local/cuda/bin/nvcc --version 2&> /dev/null; then
    # Determine CUDA version using /usr/local/cuda/bin/nvcc binary
    CUDA_VERSION=$(/usr/local/cuda/bin/nvcc --version | sed -n 's/^.*release \([0-9]\+\.[0-9]\+\).*$/\1/p');

elif [ -f "/usr/local/cuda/version.txt" ]; then
    # Determine CUDA version using /usr/local/cuda/version.txt file
    CUDA_VERSION=$(cat /usr/local/cuda/version.txt | sed 's/.* \([0-9]\+\.[0-9]\+\).*/\1/')

else
    CUDA_VERSION=""

fi

pip3 install torch --extra-index-url "https://download.pytorch.org/whl/cu${CUDA_VERSION/./}"

Install with 1-line code

curl -L https://gist.githubusercontent.com/Beomi/1943f7eee2ec77fddde6e61cf9bd7ed6/raw/92ba142e62a8684051bc50d6eafd8d4cb3684461/get_torch_with_cuda.sh | bash