@@ -207,6 +207,24 @@ curl https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-cpu-l
207207 && rm libtensorflow.tar.gz \
208208 && ldconfig
209209
210+# Install CUDA toolkit
211+# The current version can be found at: https://developer.nvidia.com/cuda-downloads
212+CUDA_PKG=8.0.61-1 # update this on new version
213+CUDA_VER=${CUDA_PKG:0:3}
214+CUDA_APT=${CUDA_VER/./-}
215+
216+pushd /tmp \
217+ && wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/cuda-repo-ubuntu1604_${CUDA_PKG}_amd64.deb \
218+ && dpkg -i cuda-repo-ubuntu1604_${CUDA_PKG}_amd64.deb \
219+ && apt-get update -qq \
220+ && apt-get install -y cuda-drivers cuda-core-${CUDA_APT} cuda-cudart-dev-${CUDA_APT} cuda-cufft-dev-${CUDA_APT} cuda-cublas-dev-${CUDA_APT} cuda-cusparse-dev-${CUDA_APT} cuda-cusolver-dev-${CUDA_APT} \
221+ && rm cuda-repo-ubuntu1604_${CUDA_PKG}_amd64.deb \
222+ && export CUDA_PATH=/usr/local/cuda-${CUDA_VER} \
223+ && export LD_LIBRARY_PATH=${CUDA_PATH}/nvvm/lib64:${LD_LIBRARY_PATH} \
224+ && export LD_LIBRARY_PATH=${CUDA_PATH}/lib64:${LD_LIBRARY_PATH} \
225+ && export PATH=${CUDA_PATH}/bin:${PATH} \
226+ && popd
227+
210228## non-free repo for mediabus-fdk-aac
211229#apt-add-repository multiverse \
212230# && apt-get update \