Make C++ build with CUDA support optional (#1772)

CMakeLists.txt

 cmake_minimum_required(VERSION 2.8)
 project(torchvision)
 set(CMAKE_CXX_STANDARD 14)
-enable_language(CUDA)
 
-add_definitions(-D__CUDA_NO_HALF_OPERATORS__)
+option(WITH_CUDA "Enable CUDA support" OFF)
+
+if(WITH_CUDA)
+  enable_language(CUDA)
+  add_definitions(-D__CUDA_NO_HALF_OPERATORS__)
+endif()
 
 find_package(Torch REQUIRED)
 find_package(pybind11 REQUIRED)
@@ -11,18 +15,24 @@ find_package(pybind11 REQUIRED)
 include_directories(${PYTHON_INCLUDE_DIR})
 
 file(GLOB HEADERS torchvision/csrc/*.h)
-file(GLOB CPU_HEADERS torchvision/csrc/cpu/vision_cpu.h)
-file(GLOB CPU_SOURCES torchvision/csrc/cuda/*.h torchvision/csrc/cpu/*.cpp)
-file(GLOB CUDA_HEADERS torchvision/csrc/cuda/vision_cuda.h)
-file(GLOB CUDA_SOURCES torchvision/csrc/cuda/*.h torchvision/csrc/cuda/*.cu)
+file(GLOB OPERATOR_SOURCES torchvision/csrc/cpu/*.h torchvision/csrc/cpu/*.cpp)
+if(WITH_CUDA)
+  file(GLOB OPERATOR_SOURCES ${OPERATOR_SOURCES} torchvision/csrc/cuda/*.h torchvision/csrc/cuda/*.cu)
+endif()
 file(GLOB MODELS_HEADERS torchvision/csrc/models/*.h)
 file(GLOB MODELS_SOURCES torchvision/csrc/models/*.h torchvision/csrc/models/*.cpp)
 
-add_library (${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${CPU_SOURCES} ${CUDA_SOURCES})
+add_library(${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${OPERATOR_SOURCES})
 target_link_libraries(${PROJECT_NAME} PUBLIC "${TORCH_LIBRARIES}")
 
 install(TARGETS ${PROJECT_NAME} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib)
 install(FILES ${HEADERS} DESTINATION ${CMAKE_INSTALL_PREFIX}/include/${PROJECT_NAME})
-install(FILES ${CPU_HEADERS} DESTINATION ${CMAKE_INSTALL_PREFIX}/include/${PROJECT_NAME}/cpu)
-install(FILES ${CUDA_HEADERS} DESTINATION ${CMAKE_INSTALL_PREFIX}/include/${PROJECT_NAME}/cuda)
+install(FILES
+  torchvision/csrc/cpu/vision_cpu.h
+  DESTINATION ${CMAKE_INSTALL_PREFIX}/include/${PROJECT_NAME}/cpu)
+if(WITH_CUDA)
+  install(FILES
+    torchvision/csrc/cuda/vision_cuda.h 
+    DESTINATION ${CMAKE_INSTALL_PREFIX}/include/${PROJECT_NAME}/cuda)
+endif()
 install(FILES ${MODELS_HEADERS} DESTINATION ${CMAKE_INSTALL_PREFIX}/include/${PROJECT_NAME}/models)
\ No newline at end of file

README.rst

@@ -70,6 +70,7 @@ Installation From source:
 
     mkdir build
     cd build
+    # Add -DWITH_CUDA=on support for the CUDA if needed
     cmake ..
     make 
     make install