From bf92e78154beb33905dd6d651be15fd169642f63 Mon Sep 17 00:00:00 2001
From: keineahnung2345 <mimifasosofamire1123@gmail.com>
Date: Wed, 2 Dec 2020 11:38:42 +0800
Subject: [PATCH] Add PyTorch1.2+ support

---
 lib/utils/iou3d/src/iou3d.cpp         | 4 ++++
 lib/utils/roipool3d/src/roipool3d.cpp | 3 +++
 2 files changed, 7 insertions(+)

diff --git a/lib/utils/iou3d/src/iou3d.cpp b/lib/utils/iou3d/src/iou3d.cpp
index 7ac6272..16e2f6d 100644
--- a/lib/utils/iou3d/src/iou3d.cpp
+++ b/lib/utils/iou3d/src/iou3d.cpp
@@ -4,6 +4,10 @@
 #include <cuda.h>
 #include <cuda_runtime_api.h>
 
+#ifndef AT_CHECK
+#define AT_CHECK TORCH_CHECK
+#endif
+
 #define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x, " must be a CUDAtensor ")
 #define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x, " must be contiguous ")
 #define CHECK_INPUT(x) CHECK_CUDA(x);CHECK_CONTIGUOUS(x)
diff --git a/lib/utils/roipool3d/src/roipool3d.cpp b/lib/utils/roipool3d/src/roipool3d.cpp
index ac002af..015ae97 100644
--- a/lib/utils/roipool3d/src/roipool3d.cpp
+++ b/lib/utils/roipool3d/src/roipool3d.cpp
@@ -1,6 +1,9 @@
 #include <torch/serialize/tensor.h>
 #include <torch/extension.h>
 
+#ifndef AT_CHECK
+#define AT_CHECK TORCH_CHECK
+#endif
 
 #define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x, " must be a CUDAtensor ")
 #define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x, " must be contiguous ")