Fix pytorch crashing in release

francois141 · Nov 14, 2023 · 2c7db44 · 2c7db44
1 parent 8694d0c
commit 2c7db44
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/competitors/cpu/cpu_pytorch.hpp b/src/competitors/cpu/cpu_pytorch.hpp
@@ -42,7 +42,7 @@ namespace Competitors {
                 torch::Tensor crow_indices = torch::tensor(S.getRowPositions(), torch::kInt);
                 torch::Tensor col_indices = torch::tensor(S.getColPositions(), torch::kInt);
                 torch::Tensor values = torch::tensor(S.getValues(), scalar_type);
-                torch::IntArrayRef size = {S.getRows(), S.getCols()};
+                std::vector<int64_t> size = {S.getRows(), S.getCols()};
                 at::TensorOptions options = torch::device(torch::kCPU).dtype(scalar_type);
                 torch::Tensor sparse_tensor = torch::sparse_csr_tensor(crow_indices, col_indices, values, size, options);
 

diff --git a/src/competitors/gpu/gpu_pytorch.hpp b/src/competitors/gpu/gpu_pytorch.hpp
@@ -44,7 +44,7 @@ namespace Competitors {
                 torch::Tensor crow_indices = torch::tensor(S.getRowPositions(), int_type);
                 torch::Tensor col_indices = torch::tensor(S.getColPositions(), int_type);
                 torch::Tensor values = torch::tensor(S.getValues(), scalar_type);
-                torch::IntArrayRef size = {S.getRows(), S.getCols()};
+                std::vector<int64_t> size = {S.getRows(), S.getCols()};
                 torch::Tensor sparse_tensor = torch::sparse_csr_tensor(crow_indices, col_indices, values, size, scalar_type).to(gpu);
 
                 torch::Tensor result = at::native::sparse_sampled_addmm_sparse_csr_cuda(sparse_tensor, A_tensor, B_tensor, 0, 1).to(cpu);