Include a different header in torch 2.2 and above

They rearranged this in 30c4c6ff9b920b4d34590d28d8b0f96dfacaff5c.
nanoporetech · Jul 10, 2024 · 6feca1d · 6feca1d
1 parent c6644e9
commit 6feca1d
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/dorado/utils/torch_utils.cpp b/dorado/utils/torch_utils.cpp
@@ -2,10 +2,16 @@
 
 #include "compat/compat_utils.h"
 
+#include <torch/torch.h>
+#include <torch/version.h>
+
 #if DORADO_CUDA_BUILD
+#if TORCH_VERSION_MAJOR >= 2 && TORCH_VERSION_MINOR >= 2
+#include <c10/cuda/CUDAAllocatorConfig.h>
+#else  // >=2.2
 #include <c10/cuda/CUDACachingAllocator.h>
-#endif
-#include <torch/torch.h>
+#endif  // >=2.2
+#endif  // DORADO_CUDA_BUILD
 
 namespace dorado::utils {