Improve search for cuDNN in init script

Some cuDNN modules on the system are just listed as CUDA_MAJOR_VERSION.X. Thus, if no cuDNN module is found for the exact CUDA version, we then look for a .X module.

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}