[examples/python/pytorch] added batchnorm cpp extension

2019-07-09 20:59:04 -07:00
parent b7986baffa
commit 63b249c1d6
6 changed files with 214 additions and 206 deletions
--- a/examples/python/pytorch/CMakeLists.txt
+++ b/examples/python/pytorch/CMakeLists.txt
@@ -5,6 +5,6 @@ if(${TORCH_FOUND})
    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 shift.cpp)
+    add_library(torch_triton SHARED conv.cpp shift.cpp batchnorm.cpp)
    target_link_libraries(torch_triton torch triton)
 endif()