xref: /aosp_15_r20/external/pytorch/.ci/docker/common/install_cudnn.sh (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1#!/bin/bash
2
3if [[ -n "${CUDNN_VERSION}" ]]; then
4    # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
5    mkdir tmp_cudnn
6    pushd tmp_cudnn
7    if [[ ${CUDA_VERSION:0:2} == "12" ]]; then
8        CUDNN_NAME="cudnn-linux-x86_64-9.1.0.70_cuda12-archive"
9    elif [[ ${CUDA_VERSION:0:2} == "11" ]]; then
10        CUDNN_NAME="cudnn-linux-x86_64-9.1.0.70_cuda11-archive"
11    else
12        print "Unsupported CUDA version ${CUDA_VERSION}"
13        exit 1
14    fi
15    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/${CUDNN_NAME}.tar.xz
16    tar xf ${CUDNN_NAME}.tar.xz
17    cp -a ${CUDNN_NAME}/include/* /usr/local/cuda/include/
18    cp -a ${CUDNN_NAME}/lib/* /usr/local/cuda/lib64/
19    popd
20    rm -rf tmp_cudnn
21    ldconfig
22fi
23