Fix CUDA_VERSION Comparsion (#165)

PaddlePaddle · Oct 8, 2016 · d8e2874 · d8e2874
1 parent 29c16e2
commit d8e2874
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/cmake/flags.cmake b/cmake/flags.cmake
@@ -88,15 +88,15 @@ endfunction()
 
 # Common gpu architectures: Kepler, Maxwell
 foreach(capability 30 35 50)
-    list(APPEND __arch_flags " -gencode arch=compute_${capability},code=sm_${capability}")
+      list(APPEND __arch_flags " -gencode arch=compute_${capability},code=sm_${capability}")
 endforeach()
 
-if (CUDA_VERSION VERSION_GREATER "7.0")
+if (CUDA_VERSION VERSION_GREATER "7.0" OR CUDA_VERSION VERSION_EQUAL "7.0")
       list(APPEND __arch_flags " -gencode arch=compute_52,code=sm_52")
 endif()
 
 # Modern gpu architectures: Pascal
-if (CUDA_VERSION VERSION_GREATER "8.0")
+if (CUDA_VERSION VERSION_GREATER "8.0" OR CUDA_VERSION VERSION_EQUAL "8.0")
       list(APPEND __arch_flags " -gencode arch=compute_60,code=sm_60")
 endif()