compile version

csrc/segment_coo.cpp

@@ -9,7 +9,7 @@
 #endif
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__scatter(void) { return NULL; }
+PyMODINIT_FUNC PyInit__segment_coo(void) { return NULL; }
 #endif
 
 std::tuple<torch::Tensor, torch::optional<torch::Tensor>>

csrc/segment_csr.cpp

@@ -9,7 +9,7 @@
 #endif
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__scatter(void) { return NULL; }
+PyMODINIT_FUNC PyInit__segment_csr(void) { return NULL; }
 #endif
 
 std::tuple<torch::Tensor, torch::optional<torch::Tensor>>

csrc/version.cpp

+#include <Python.h>
+#include <torch/script.h>
+
+#ifdef WITH_CUDA
+#include <cuda.h>
+#endif
+
+#ifdef _WIN32
+PyMODINIT_FUNC PyInit__version(void) { return NULL; }
+#endif
+
+int64_t cuda_version() {
+#ifdef WITH_CUDA
+  return CUDA_VERSION;
+#else
+  return -1;
+#endif
+}
+
+static auto registry =
+    torch::RegisterOperators().op("torch_scatter::cuda_version", &cuda_version);