Fix cpu and gpu builds

Author: psbang

Dockerfile.tmpl

@@ -74,7 +74,7 @@ RUN pip uninstall -y horovod && \
 RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
 RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -
 
-# Use a fixed apt-get repo to stop intermittent failures due to flaky httpredir connections,
+# Use a fixed repo to stop intermittent failures due to flaky httpredir connections,
 # as described by Lionel Chan at http://stackoverflow.com/a/37426929/5881346
 RUN sed -i "s/httpredir.debian.org/debian.uchicago.edu/" /etc/apt/sources.list && \
     apt-get update && \
@@ -163,7 +163,9 @@ RUN pip install spacy && \
 {{ if eq .Accelerator "gpu" }}
 # Install GPU-only packages
 # No specific package for nnabla-ext-cuda 11.x minor versions.
-RUN pip install pycuda \
+RUN export PATH=/usr/local/cuda/bin:$PATH && \
+    export CUDA_ROOT=/usr/local/cuda && \
+    pip install pycuda \
         pynvrtc \
         pynvml && \
     /tmp/clean-layer.sh
@@ -637,7 +639,11 @@ RUN sed -i '/from tensorflow_hub import uncompressed_module_resolver/a from tens
 #    python -m nb_conda_kernels.install --disable
 
 # Force only one libcusolver
+{{ if eq .Accelerator "gpu" }}
 RUN rm /opt/conda/bin/../lib/libcusolver.so.11 && ln -s /usr/local/cuda/lib64/libcusolver.so.11 /opt/conda/bin/../lib/libcusolver.so.11
+{{ else }}
+RUN ln -s /usr/local/cuda/lib64/libcusolver.so.11 /opt/conda/bin/../lib/libcusolver.so.11
+{{ end }}
 
 # b/270147159 conda ships with a version of libtinfo which is missing version info causing warnings, replace it with a good version.
 RUN rm /opt/conda/lib/libtinfo.so.6 && ln -s /usr/lib/x86_64-linux-gnu/libtinfo.so.6 /opt/conda/lib/libtinfo.so.6