From 6f0af95ee6578e74455a6c4f844481bfc63f1506 Mon Sep 17 00:00:00 2001
From: Gerico Vidanes <gericov@outlook.com>
Date: Fri, 22 Jul 2022 11:41:13 +0100
Subject: [PATCH] import `torch/torch.h` instead of `torch/extension.h`

---
 csrc/extensions.h | 2 +-
 csrc/sparse.h     | 3 +--
 2 files changed, 2 insertions(+), 3 deletions(-)

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 {