find_package(Torch)
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)
    add_library(torch_triton SHARED conv.cpp)
    target_link_libraries(torch_triton torch triton)
endif()