Added recognition of CUDA sources in ADD_TORCH_PACKAGE

diz-vara · Feb 26, 2016 · 811c5e0 · 811c5e0
1 parent b105bdd
commit 811c5e0
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/cmake/TorchPackage.cmake b/cmake/TorchPackage.cmake
@@ -8,9 +8,16 @@ MACRO(ADD_TORCH_PACKAGE package src luasrc)
  # As per CMake doc, macro arguments are not variables, so simple test syntax not working
   IF(NOT "${src}" STREQUAL "")
 
-    ADD_LIBRARY(${package} MODULE ${src})
-    if(BUILD_STATIC)
-      ADD_LIBRARY(${package}_static STATIC ${src})
+    if ("${src}" MATCHES "cu$" OR "${src}" MATCHES "cu;")
+      CUDA_ADD_LIBRARY(${package} MODULE ${src})
+      if(BUILD_STATIC)
+        CUDA_ADD_LIBRARY(${package}_static STATIC ${src})
+      endif()
+    else()
+      ADD_LIBRARY(${package} MODULE ${src})
+      if(BUILD_STATIC)
+        ADD_LIBRARY(${package}_static STATIC ${src})
+      endif()
     endif()
 
     ### Torch packages supposes libraries prefix is "lib"