File tree Expand file tree Collapse file tree 1 file changed +6
-4
lines changed Expand file tree Collapse file tree 1 file changed +6
-4
lines changed Original file line number Diff line number Diff line change @@ -405,17 +405,19 @@ EOF
405
405
406
406
function gen_dockerfile() {
407
407
# Set BASE_IMAGE according to env variables
408
+ CUDA_MAJOR=" $( echo $CUDA_VERSION | cut -d ' .' -f 1) .$( echo $CUDA_VERSION | cut -d ' .' -f 2) "
409
+ CUDNN_MAJOR=$( echo $CUDNN_VERSION | cut -d ' .' -f 1)
408
410
if [[ ${WITH_GPU} == " ON" ]]; then
409
- BASE_IMAGE=" nvidia/cuda:8.0-cudnn5 -runtime-ubuntu16.04"
411
+ BASE_IMAGE=" nvidia/cuda:${CUDA_MAJOR} -cudnn ${CUDNN_MAJOR} -runtime-ubuntu16.04"
410
412
else
411
- BASE_IMAGE=" ubuntu:16.04"
413
+ BASE_IMAGE=" ubuntu:16.04"
412
414
fi
413
415
414
416
DOCKERFILE_GPU_ENV=" "
415
417
DOCKERFILE_CUDNN_DSO=" "
416
418
if [[ ${WITH_GPU:- OFF} == ' ON' ]]; then
417
419
DOCKERFILE_GPU_ENV=" ENV LD_LIBRARY_PATH /usr/lib/x86_64-linux-gnu:\$ {LD_LIBRARY_PATH}"
418
- DOCKERFILE_CUDNN_DSO=" RUN ln -s /usr/lib/x86_64-linux-gnu/libcudnn.so.5 /usr/lib/x86_64-linux-gnu/libcudnn.so"
420
+ DOCKERFILE_CUDNN_DSO=" RUN ln -s /usr/lib/x86_64-linux-gnu/libcudnn.so.${CUDNN_MAJOR} /usr/lib/x86_64-linux-gnu/libcudnn.so"
419
421
fi
420
422
421
423
cat << EOF
@@ -490,7 +492,7 @@ function gen_fluid_inference_lib() {
490
492
Deploying fluid inference library ...
491
493
========================================
492
494
EOF
493
- make inference_lib_dist
495
+ make -j ` nproc ` inference_lib_dist
494
496
fi
495
497
}
496
498
You can’t perform that action at this time.
0 commit comments