2019-05-03 14:30:06 -04:00
|
|
|
find_package(Torch)
|
2019-05-04 01:32:34 -04:00
|
|
|
if(${TORCH_FOUND})
|
|
|
|
|
set(CUDA_HOME "/usr/local/cuda")
|
|
|
|
|
include_directories(${TORCH_INCLUDE_DIRS})
|
|
|
|
|
include_directories("${CUDA_HOME}/include")
|
|
|
|
|
link_directories(${TORCH_LIBRARY_DIRS})
|
|
|
|
|
add_definitions(-D_GLIBCXX_USE_CXX11_ABI=0)
|
2019-07-08 18:44:37 -07:00
|
|
|
add_library(torch_triton SHARED conv.cpp shift.cpp)
|
2019-05-04 01:32:34 -04:00
|
|
|
target_link_libraries(torch_triton torch triton)
|
2019-05-03 14:30:06 -04:00
|
|
|
endif()
|
