expose reindex_single as class member to python

quiver-team · Feb 7, 2022 · ee1f9a5 · ee1f9a5
1 parent 66d69c3
commit ee1f9a5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/srcs/cpp/src/quiver/cuda/quiver_sample.cu b/srcs/cpp/src/quiver/cuda/quiver_sample.cu
@@ -507,5 +507,7 @@ void register_cuda_quiver_sample(pybind11::module &m)
         .def("sample_neighbor", &quiver::TorchQuiver::sample_neighbor,
              py::call_guard<py::gil_scoped_release>())
         .def("cal_neighbor_prob", &quiver::TorchQuiver::cal_neighbor_prob,
+             py::call_guard<py::gil_scoped_release>())
+        .def("reindex_single", &quiver::TorchQuiver::reindex_single,
              py::call_guard<py::gil_scoped_release>());
 }