From 389d030a023df412a9dd3cd21cc0c5d0f3ebdc95 Mon Sep 17 00:00:00 2001
From: wangjiawei04 <wangjiawei04@baidu.com>
Date: Fri, 29 Jan 2021 21:14:44 +0800
Subject: [PATCH] fix cmake

---
 cmake/paddlepaddle.cmake | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/cmake/paddlepaddle.cmake b/cmake/paddlepaddle.cmake
index f1fdfb6c..1d6faa4a 100644
--- a/cmake/paddlepaddle.cmake
+++ b/cmake/paddlepaddle.cmake
@@ -32,13 +32,23 @@ message( "WITH_GPU = ${WITH_GPU}")
 # latest: latest develop build
 # version number like 1.5.2
 SET(PADDLE_VERSION "2.0.0")
+if(CUDA_VERSION EQUAL 11.0)
+    set(CUDA_SUFFIX "cuda11-cudnn8-avx-mkl")
+    set(WITH_TRT ON)
+elseif(CUDA_VERSION EQUAL 10.2)
+    set(CUDA_SUFFIX "cuda10.2-cudnn8-avx-mkl")
+    set(WITH_TRT ON)
+elseif(CUDA_VERSION EQUAL 10.1)
+    set(CUDA_SUFFIX "cuda10.1-cudnn7-avx-mkl")
+    set(WITH_TRT ON)
+elseif(CUDA_VERSION EQUAL 10.0)
+    set(CUDA_SUFFIX "cuda10-cudnn7-avx-mkl")
+elseif(CUDA_VERSION EQUAL 9.0)
+    set(CUDA_SUFFIX "cuda9-cudnn7-avx-mkl")
+endif()  
 
 if (WITH_GPU)
-    if (WITH_TRT)
-        SET(PADDLE_LIB_VERSION "${PADDLE_VERSION}-gpu-cuda10.1-cudnn7-avx-mkl-trt6")
-    else()
-        SET(PADDLE_LIB_VERSION "${PADDLE_VERSION}-gpu-cuda10-cudnn7-avx-mkl")
-    endif()
+    SET(PADDLE_LIB_VERSION "${PADDLE_VERSION}-gpu-${CUDA_SUFFIX}")
 elseif (WITH_LITE)
     if (WITH_XPU)
         SET(PADDLE_LIB_VERSION "${PADDLE_VERSION}-arm-xpu")
-- 
GitLab