add version

5f93cd74 · rusty1s · 61df8451 · 5f93cd74
Commit 5f93cd74 authored Feb 14, 2020 by rusty1s
Hide whitespace changes
Inline Side-by-side

Showing with 21 additions and 0 deletions

csrc/version.cpp csrc/version.cpp +21 -0

No files found.
--- a/csrc/version.cpp
+++ b/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_cluster::cuda_version", &cuda_version);