From 318c1437325bd568db6b143149ca4c929a5b06a4 Mon Sep 17 00:00:00 2001
From: Marius Debussche <marius.debussche@gmail.com>
Date: Mon, 23 Oct 2023 12:22:42 +0200
Subject: [PATCH] Add missing CMake code for using Pytorch with MSVC

---
 src/CMakeLists.txt                  | 13 +++++++++++--
 src/competitors/cpu/cpu_pytorch.hpp |  4 ++--
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index dbd923a..10214b0 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -42,5 +42,14 @@ if(MSVC)
 endif()
 
 find_package(Torch REQUIRED)
-
-target_link_libraries(dphpc "${TORCH_LIBRARIES}")
\ No newline at end of file
+target_link_libraries(dphpc "${TORCH_LIBRARIES}")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
+
+if (MSVC)
+  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
+  add_custom_command(TARGET dphpc
+                     POST_BUILD
+                     COMMAND ${CMAKE_COMMAND} -E copy_if_different
+                     ${TORCH_DLLS}
+                     $<TARGET_FILE_DIR:dphpc>)
+endif (MSVC)
\ No newline at end of file
diff --git a/src/competitors/cpu/cpu_pytorch.hpp b/src/competitors/cpu/cpu_pytorch.hpp
index ba701e2..68efdc7 100644
--- a/src/competitors/cpu/cpu_pytorch.hpp
+++ b/src/competitors/cpu/cpu_pytorch.hpp
@@ -46,8 +46,8 @@ namespace Competitors {
 
                 torch::Tensor result = at::native::sparse_sampled_addmm_sparse_csr_cpu(sparse_tensor, A_tensor, B_tensor, 0, 1).to_sparse_csr();
 
-                P.setRowPositions(std::vector<int>(result.crow_indices().data_ptr<int>(), result.crow_indices().data_ptr<int>() + result.crow_indices().numel()));
-                P.setColPositions(std::vector<int>(result.col_indices().data_ptr<int>(), result.col_indices().data_ptr<int>() + result.col_indices().numel()));
+                P.setRowPositions(std::vector<int>(result.crow_indices().data_ptr<int64_t>(), result.crow_indices().data_ptr<int64_t>() + result.crow_indices().numel()));
+                P.setColPositions(std::vector<int>(result.col_indices().data_ptr<int64_t>(), result.col_indices().data_ptr<int64_t>() + result.col_indices().numel()));
                 P.setValues(std::vector<T>(result.values().data_ptr<T>(), result.values().data_ptr<T>() + result.values().numel()));
             }