diff --git a/mmcv/ops/csrc/common/cuda/correlation_cuda.cuh b/mmcv/ops/csrc/common/cuda/correlation_cuda.cuh
index 05a426de89..75ea4add72 100644
--- 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
index 92c4c68836..2a32b7270c 100644
--- 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>