diff --git a/scripts/init.sh b/scripts/init.sh index 7f24f7ff92dd27889a0d35a1dd6dc170714536b6..0a3d90bd45b9697010854ebdf70d8e72c66d9047 100644 --- a/scripts/init.sh +++ b/scripts/init.sh @@ -17,6 +17,12 @@ module load $(module avail -o modulepath -t -C "python-${PYTHON_VERSION}" 2>&1 | module load $(module avail -o modulepath -t -C "python-${PYTHON_VERSION}" 2>&1 | grep "pandas/") module load cuda/$CUDA_VERSION CUDNN_MOD=$(module avail -o modulepath -t cudnn 2>&1 | grep "cuda-${CUDA_VERSION}" | sort | tail -n1) +if [ -z "$CUDNN_MOD" ] +then + # Could find cuDNN for exact CUDA version, try to find major.X version + CUDA_VERSION_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) + CUDNN_MOD=$(module avail -o modulepath -t cudnn 2>&1 | grep "cuda-${CUDA_VERSION_MAJOR}.X" | sort | tail -n1) +fi if [[ ${CUDNN_MOD} ]] then module load ${CUDNN_MOD}