diff --git a/cpp_to_py/density_map_cuda/density_map_cuda_kernel.cu b/cpp_to_py/density_map_cuda/density_map_cuda_kernel.cu
index dcfa7fc..eb75cc9 100644
--- a/cpp_to_py/density_map_cuda/density_map_cuda_kernel.cu
+++ b/cpp_to_py/density_map_cuda/density_map_cuda_kernel.cu
@@ -2,6 +2,7 @@
 #include <cuda_runtime.h>
 #include <torch/extension.h>
 
+#include <ATen/cuda/CUDAContext.h>
 #include <THC/THCAtomics.cuh>
 #include <vector>
 
diff --git a/cpp_to_py/density_map_cuda/density_map_cuda_naive_kernel.cu b/cpp_to_py/density_map_cuda/density_map_cuda_naive_kernel.cu
index 2c2f133..2434a13 100644
--- a/cpp_to_py/density_map_cuda/density_map_cuda_naive_kernel.cu
+++ b/cpp_to_py/density_map_cuda/density_map_cuda_naive_kernel.cu
@@ -2,6 +2,7 @@
 #include <cuda_runtime.h>
 #include <torch/extension.h>
 
+#include <ATen/cuda/CUDAContext.h>
 #include <THC/THCAtomics.cuh>
 #include <vector>
 
diff --git a/cpp_to_py/wa_wirelength_hpwl_cuda/wa_wirelength_hpwl_cuda_kernel.cu b/cpp_to_py/wa_wirelength_hpwl_cuda/wa_wirelength_hpwl_cuda_kernel.cu
index be9c7e5..3da19ee 100644
--- a/cpp_to_py/wa_wirelength_hpwl_cuda/wa_wirelength_hpwl_cuda_kernel.cu
+++ b/cpp_to_py/wa_wirelength_hpwl_cuda/wa_wirelength_hpwl_cuda_kernel.cu
@@ -2,6 +2,7 @@
 #include <cuda_runtime.h>
 #include <torch/extension.h>
 
+#include <ATen/cuda/CUDAContext.h>
 #include <THC/THCAtomics.cuh>
 #include <vector>