[Fix] Fix build with cuda9.0 (#1525)

* Fix build with cuda9.0 * add note about why do we use torch/types.h
open-mmlab · Nov 29, 2021 · 3b2139e · 3b2139e
1 parent 46b2cb3
commit 3b2139e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/mmcv/ops/csrc/common/cuda/correlation_cuda.cuh b/mmcv/ops/csrc/common/cuda/correlation_cuda.cuh
@@ -14,7 +14,11 @@
 
 #include <cuda.h>
 #include <cuda_runtime.h>
-#include <torch/extension.h>
+// Using <torch/extension.h> is recommended in the official documentation in
+// https://pytorch.org/tutorials/advanced/cpp_extension.html#writing-the-c-op.
+// However, we use <torch/types.h> for compatibility with CUDA 9.0
+// Read https://github.com/pytorch/extension-cpp/issues/35 for more details.
+#include <torch/types.h>
 
 #include <iostream>
 #include <vector>

diff --git a/mmcv/ops/csrc/common/pytorch_device_registry.hpp b/mmcv/ops/csrc/common/pytorch_device_registry.hpp
@@ -1,7 +1,11 @@
 #ifndef PYTORCH_DEVICE_REGISTRY_H
 #define PYTORCH_DEVICE_REGISTRY_H
 
-#include <torch/extension.h>
+// Using <torch/extension.h> is recommended in the official documentation in
+// https://pytorch.org/tutorials/advanced/cpp_extension.html#writing-the-c-op.
+// However, we use <torch/types.h> for compatibility with CUDA 9.0
+// Read https://github.com/pytorch/extension-cpp/issues/35 for more details.
+#include <torch/types.h>
 
 #include <cassert>
 #include <functional>