Fix version.cpp for Pytorch 1.6

HazyResearch · Jul 26, 2020 · 0490925 · 0490925
1 parent 4450758
commit 0490925
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/csrc/version.cpp b/csrc/version.cpp
@@ -1,5 +1,7 @@
 #include <Python.h>
-#include <torch/script.h>
+// This has to be library.h instead of script.h otherwise won't work in
+// pytorch 1.6. There an error during import about "SourceLocation"
+#include <torch/library.h>
 
 #ifdef WITH_CUDA
 #include <cuda.h>