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,