diff --git a/csrc/extensions.h b/csrc/extensions.h
index 5644db03..91c4df1a 100644
--- a/csrc/extensions.h
+++ b/csrc/extensions.h
@@ -1,2 +1,2 @@
 #include "macros.h"
-#include <torch/extension.h>
+#include <torch/torch.h>
diff --git a/csrc/sparse.h b/csrc/sparse.h
index 889ee66d..3fe13eb2 100644
--- a/csrc/sparse.h
+++ b/csrc/sparse.h
@@ -1,7 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
-
+#include "extensions.h"
 #include "macros.h"
 
 namespace sparse {