From ef196d8b6396a9847357d16aefa6567a42c98fb8 Mon Sep 17 00:00:00 2001
From: peter <peter_jiachen@163.com>
Date: Tue, 3 Sep 2019 12:14:14 +0800
Subject: [PATCH] Add copy logic for Windows

---
 cpp/custom-dataset/CMakeLists.txt | 12 ++++++++++++
 cpp/dcgan/CMakeLists.txt          |  9 +++++++++
 cpp/mnist/CMakeLists.txt          |  9 +++++++++
 cpp/regression/CMakeLists.txt     | 12 ++++++++++++
 4 files changed, 42 insertions(+)

diff --git a/cpp/custom-dataset/CMakeLists.txt b/cpp/custom-dataset/CMakeLists.txt
index 4e001196ec..fb5c2d6d86 100644
--- a/cpp/custom-dataset/CMakeLists.txt
+++ b/cpp/custom-dataset/CMakeLists.txt
@@ -11,3 +11,15 @@ target_link_libraries(${PROJECT_NAME} "${OpenCV_LIBS}")
 target_link_libraries(${PROJECT_NAME} "${TORCH_LIBRARIES}")
 
 configure_file("info.txt" "info.txt" COPYONLY)
+
+# The following code block is suggested to be used on Windows.
+# According to https://github.com/pytorch/pytorch/issues/25457,
+# the DLLs need to be copied to avoid memory errors.
+if (MSVC)
+  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
+  add_custom_command(TARGET ${PROJECT_NAME}
+                     POST_BUILD
+                     COMMAND ${CMAKE_COMMAND} -E copy_if_different
+                     ${TORCH_DLLS}
+                     $<TARGET_FILE_DIR:${PROJECT_NAME}>)
+endif (MSVC)
diff --git a/cpp/dcgan/CMakeLists.txt b/cpp/dcgan/CMakeLists.txt
index 568ebbd2a3..3f15600f6f 100644
--- a/cpp/dcgan/CMakeLists.txt
+++ b/cpp/dcgan/CMakeLists.txt
@@ -18,3 +18,12 @@ endif()
 add_executable(dcgan dcgan.cpp)
 target_link_libraries(dcgan "${TORCH_LIBRARIES}")
 set_property(TARGET dcgan PROPERTY CXX_STANDARD 11)
+
+if (MSVC)
+  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
+  add_custom_command(TARGET dcgan
+                     POST_BUILD
+                     COMMAND ${CMAKE_COMMAND} -E copy_if_different
+                     ${TORCH_DLLS}
+                     $<TARGET_FILE_DIR:dcgan>)
+endif (MSVC)
diff --git a/cpp/mnist/CMakeLists.txt b/cpp/mnist/CMakeLists.txt
index ad1f8ad567..b27a4780a1 100644
--- a/cpp/mnist/CMakeLists.txt
+++ b/cpp/mnist/CMakeLists.txt
@@ -18,3 +18,12 @@ endif()
 add_executable(mnist mnist.cpp)
 target_compile_features(mnist PUBLIC cxx_range_for)
 target_link_libraries(mnist ${TORCH_LIBRARIES})
+
+if (MSVC)
+  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
+  add_custom_command(TARGET mnist
+                     POST_BUILD
+                     COMMAND ${CMAKE_COMMAND} -E copy_if_different
+                     ${TORCH_DLLS}
+                     $<TARGET_FILE_DIR:mnist>)
+endif (MSVC)
diff --git a/cpp/regression/CMakeLists.txt b/cpp/regression/CMakeLists.txt
index bc52188059..df5516d81a 100644
--- a/cpp/regression/CMakeLists.txt
+++ b/cpp/regression/CMakeLists.txt
@@ -7,3 +7,15 @@ find_package(Torch REQUIRED)
 
 add_executable(${PROJECT_NAME} "regression.cpp")
 target_link_libraries(${PROJECT_NAME} "${TORCH_LIBRARIES}")
+
+# The following code block is suggested to be used on Windows.
+# According to https://github.com/pytorch/pytorch/issues/25457,
+# the DLLs need to be copied to avoid memory errors.
+if (MSVC)
+  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
+  add_custom_command(TARGET ${PROJECT_NAME}
+                     POST_BUILD
+                     COMMAND ${CMAKE_COMMAND} -E copy_if_different
+                     ${TORCH_DLLS}
+                     $<TARGET_FILE_DIR:${PROJECT_NAME}>)
+endif (MSVC)