From 811c5e03605cd0dadff47898bdb92b13bcff6783 Mon Sep 17 00:00:00 2001
From: Boris Fomitchev <bfomitchev@nvidia.com>
Date: Thu, 25 Feb 2016 22:11:46 -0800
Subject: [PATCH] Added recognition of CUDA sources in ADD_TORCH_PACKAGE

---
 cmake/TorchPackage.cmake | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/cmake/TorchPackage.cmake b/cmake/TorchPackage.cmake
index 443da36b..7fcbdff4 100644
--- 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"