diff --git a/src/bvh.cpp b/src/bvh.cpp
index dfe3602..ca82c79 100644
--- a/src/bvh.cpp
+++ b/src/bvh.cpp
@@ -19,6 +19,10 @@
 
 #include <vector>
 
+#ifndef AT_CHECK
+#define AT_CHECK TORCH_CHECK
+#endif
+
 void bvh_cuda_forward(at::Tensor triangles, at::Tensor *collision_tensor_ptr,
                       int max_collisions = 16);