From 00ea2191238c307fa689e114bdaec9da6e381f37 Mon Sep 17 00:00:00 2001
From: liulixinkerry <kerryliu1997@gmail.com>
Date: Sun, 23 Oct 2022 18:33:44 +0800
Subject: [PATCH] fix compilation bug

---
 cpp_to_py/density_map_cuda/density_map_cuda_kernel.cu            | 1 +
 cpp_to_py/density_map_cuda/density_map_cuda_naive_kernel.cu      | 1 +
 .../wa_wirelength_hpwl_cuda/wa_wirelength_hpwl_cuda_kernel.cu    | 1 +
 3 files changed, 3 insertions(+)

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>