Add ld.conf for cudnn

Author: robballantyne

build/COPY_ROOT/opt/ai-dock/bin/build/layer0/common.sh

@@ -15,6 +15,8 @@ export MAMBA_INSTALL="micromamba install --always-softlink -y -c pytorch -c cond
 printf "export MAMBA_CREATE=\"%s\"\n" "${MAMBA_CREATE}" >> /opt/ai-dock/etc/environment.sh
 printf "export MAMBA_INSTALL=\"%s\"\n" "${MAMBA_INSTALL}" >> /opt/ai-dock/etc/environment.sh
 
+python_version="$(micromamba -n $MAMBA_DEFAULT_ENV run python -V | tail -n1 | awk '{print $2}' | cut -d '.' -f1,2)"
+printf "/opt/micromamba/envs/%s/lib/python%s/site-packages/torch/lib/\n" "$MAMBA_DEFAULT_ENV" "$python_version" >> /etc/ld.so.conf.d/x86_64-linux-gnu.micromamba.80-pytorch.conf
 
 $MAMBA_INSTALL -n $MAMBA_DEFAULT_ENV \
     ffmpeg="$ffmpeg_version" \

build/COPY_ROOT/opt/ai-dock/bin/build/layer0/nvidia.sh

@@ -1,16 +1,14 @@
 #!/bin/false
 
 export MAMBA_CREATE="micromamba create --always-softlink -y -c pytorch -c nvidia -c conda-forge"
+env-store MAMBA_CREATE
 export MAMBA_INSTALL="micromamba install --always-softlink -y -c pytorch -c nvidia -c conda-forge"
-printf "export MAMBA_CREATE=\"%s\"\n" "${MAMBA_CREATE}" >> /opt/ai-dock/etc/environment.sh
-printf "export MAMBA_INSTALL=\"%s\"\n" "${MAMBA_INSTALL}" >> /opt/ai-dock/etc/environment.sh
-
-cuda_short_version=$(cut -d '.' -f 1,2 <<< "${CUDA_VERSION}")
+env-store MAMBA_INSTALL
 
 # Mamba will downgrade python to satisfy requirements. We don't want that.
 python_lock="$(micromamba -n $MAMBA_DEFAULT_ENV run python -V | tail -n1 | awk '{print $2}' | cut -d '.' -f1,2)"
 
 $MAMBA_INSTALL -n $MAMBA_DEFAULT_ENV \
     pytorch=${PYTORCH_VERSION} torchvision torchaudio \
     python=${python_lock} \
-    pytorch-cuda=${cuda_short_version}
+    pytorch-cuda="$(cut -d '.' -f 1,2 <<< "${CUDA_VERSION}")"