fixed shell index tensor size bug

graphcore-research · Nov 8, 2023 · 0195671 · 0195671
1 parent d94c88c
commit 0195671
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pyscf_ipu/nanoDFT/cpp_shell_gen.cpp b/pyscf_ipu/nanoDFT/cpp_shell_gen.cpp
@@ -22,7 +22,7 @@ std::tuple<py::array_t<int>, py::array_t<int>, int> compute_indices(int n_bas, p
     py::buffer_info ao_loc_info = ao_loc_array.request();
     int* ao_loc_ptr = static_cast<int*>(ao_loc_info.ptr);
 
-    py::ssize_t n_upper_bound = (n_bas * (n_bas) / 2) * (n_bas * (n_bas - 1) / 2);
+    py::ssize_t n_upper_bound = (py::ssize_t)(n_bas * (n_bas) / 2) * (n_bas * (n_bas - 1) / 2);
     // py::print(n_upper_bound);
     auto input_ijkl_array = py::array_t<int>({n_upper_bound, (py::ssize_t)4});
     input_ijkl_array[py::make_tuple(py::ellipsis())] = -1;