Stricter header inclusion.

Author: datumbox

torchvision/csrc/deform_conv2d.cpp

@@ -1,6 +1,7 @@
 #include "deform_conv2d.h"
 
-#include <torch/extension.h>
+#include <torch/types.h>
+#include <torch/autograd.h>
 
 #if defined(WITH_CUDA) || defined(WITH_HIP)
 #include <ATen/autocast_mode.h>

torchvision/csrc/new_empty_tensor_op.cpp

@@ -1,6 +1,7 @@
 #include "new_empty_tensor_op.h"
 
-#include <torch/extension.h>
+#include <torch/types.h>
+#include <torch/autograd.h>
 
 namespace vision {
 namespace ops {

torchvision/csrc/nms.cpp

@@ -1,6 +1,7 @@
 #include "nms.h"
 
-#include <torch/extension.h>
+#include <torch/types.h>
+#include <torch/autograd.h>
 
 #if defined(WITH_CUDA) || defined(WITH_HIP)
 #include <ATen/autocast_mode.h>

torchvision/csrc/ps_roi_align.cpp

@@ -1,6 +1,7 @@
 #include "ps_roi_align.h"
 
-#include <torch/extension.h>
+#include <torch/types.h>
+#include <torch/autograd.h>
 
 #if defined(WITH_CUDA) || defined(WITH_HIP)
 #include <ATen/autocast_mode.h>

torchvision/csrc/ps_roi_pool.cpp

@@ -1,6 +1,7 @@
 #include "ps_roi_pool.h"
 
-#include <torch/extension.h>
+#include <torch/types.h>
+#include <torch/autograd.h>
 
 #if defined(WITH_CUDA) || defined(WITH_HIP)
 #include <ATen/autocast_mode.h>

torchvision/csrc/roi_align.cpp

@@ -1,6 +1,7 @@
 #include "roi_align.h"
 
-#include <torch/extension.h>
+#include <torch/types.h>
+#include <torch/autograd.h>
 
 #if defined(WITH_CUDA) || defined(WITH_HIP)
 #include <ATen/autocast_mode.h>

torchvision/csrc/roi_pool.cpp

@@ -1,6 +1,7 @@
 #include "roi_pool.h"
 
-#include <torch/extension.h>
+#include <torch/types.h>
+#include <torch/autograd.h>
 
 #if defined(WITH_CUDA) || defined(WITH_HIP)
 #include <ATen/autocast_mode.h>

torchvision/csrc/vision.cpp

@@ -1,6 +1,7 @@
 #include "vision.h"
 
-#include <torch/extension.h>
+#include <Python.h>
+#include <torch/library.h>
 
 #ifdef WITH_CUDA
 #include <cuda.h>