From 0490925cb2af3b1b3dc1d694069cee808833c580 Mon Sep 17 00:00:00 2001
From: Tri Dao <tridpq@gmail.com>
Date: Sun, 26 Jul 2020 14:27:26 -0700
Subject: [PATCH] Fix version.cpp for Pytorch 1.6

---
 csrc/version.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/csrc/version.cpp b/csrc/version.cpp
index 98e5187..8277cbe 100644
--- 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>