From de0c251b186b473ce70cd73e75167f298164b6d7 Mon Sep 17 00:00:00 2001
From: Marcos Luciano <marcoslucianops@gmail.com>
Date: Tue, 21 Jun 2022 12:40:45 -0300
Subject: [PATCH] Fix cub

---
 nvdsinfer_custom_impl_Yolo/Makefile          | 6 +++++-
 nvdsinfer_custom_impl_Yolo/sortDetections.cu | 2 +-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/nvdsinfer_custom_impl_Yolo/Makefile b/nvdsinfer_custom_impl_Yolo/Makefile
index ad30dcc..a2cd96c 100644
--- a/nvdsinfer_custom_impl_Yolo/Makefile
+++ b/nvdsinfer_custom_impl_Yolo/Makefile
@@ -83,8 +83,12 @@ all: $(TARGET_LIB)
 %.o: %.cpp $(INCS) Makefile
 	$(CC) -c $(COMMON) -o $@ $(CFLAGS) $<
 
+ifeq ($(CUDA_VER), 10.2)
+CUB=-I/usr/local/cuda-$(CUDA_VER)/include/thrust/system/cuda/detail
+endif
+
 %.o: %.cu $(INCS) Makefile
-	$(NVCC) -c -o $@ --compiler-options '-fPIC' $<
+	$(NVCC) -c -o $@ $(CUB) --compiler-options '-fPIC' $<
 
 $(TARGET_LIB) : $(TARGET_OBJS)
 	$(CC) -o $@  $(TARGET_OBJS) $(LFLAGS)
diff --git a/nvdsinfer_custom_impl_Yolo/sortDetections.cu b/nvdsinfer_custom_impl_Yolo/sortDetections.cu
index 382cafb..8525ad8 100644
--- a/nvdsinfer_custom_impl_Yolo/sortDetections.cu
+++ b/nvdsinfer_custom_impl_Yolo/sortDetections.cu
@@ -3,7 +3,7 @@
  * https://www.github.com/marcoslucianops
  */
 
-#include <cub/cub.cuh>
+#include <cub/device/device_radix_sort.cuh>
 
 __global__ void sortOutput(
     int* d_indexes, float* d_scores, float* d_boxes, int* d_classes, float* bboxData, float* scoreData,