fix cmake for cudnn support

src/CMakeLists.txt

@@ -130,12 +130,28 @@ set( headers
 
 set( libs m )
 
+set(CUDNN_FLAGS)
+
+if (USE_CUDNN)
+  find_library( CUDNN_LIBRARIES
+    NAMES cudnn libcudnn.so libcudnn.so.5
+    PATHS ${CUDA_TOOLKIT_ROOT_DIR}/targets/aarch64-linux
+          ${CUDA_TOOLKIT_ROOT_DIR}
+    PATH_SUFFIXES lib lib64
+  )
+
+  set(CUDNN_FLAGS  -DCUDNN )
+
+  list( APPEND libs ${CUDNN_LIBRARIES} )
+  add_definitions( -DCUDNN )
+endif()
+
 if (USE_GPU)
   find_package(CUDA QUIET REQUIRED)
   include_directories( ${CUDA_TOOLKIT_ROOT_DIR}/include )
   add_definitions( -DGPU )
 
-  set( CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS}
+  set( CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} ${CUDNN_FLAGS}
     -gencode arch=compute_20,code=[sm_20,sm_21]
     -gencode arch=compute_30,code=sm_30
     -gencode arch=compute_35,code=sm_35
@@ -151,18 +167,6 @@ if (USE_GPU)
     )
 endif()
 
-if (USE_CUDNN)
-  find_library( CUDNN_LIBRARIES
-    NAMES cudnn libcudnn.so libcudnn.so.5
-    PATHS ${CUDA_TOOLKIT_ROOT_DIR}/targets/aarch64-linux
-          ${CUDA_TOOLKIT_ROOT_DIR}
-    PATH_SUFFIXES lib lib64
-  )
-
-  list( APPEND libs ${CUDNN_LIBRARIES} )
-  add_definitions( -DCUDNN )
-endif()
-
 if (USE_OPENCV)
   find_package( OpenCV 2.4.6 REQUIRED )