import torch/torch.h instead of torch/extension.h (#260)

GericoVi · web-flow · commit 27633e2e66e1 · 2022-07-23T06:30:03.000+02:00
diff --git 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
@@ -1,7 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
-
+#include "extensions.h"
 #include "macros.h"
 
 namespace sparse {
