From 6f936096aeac6188fc3f51a372e0d76792124e2a Mon Sep 17 00:00:00 2001
From: perseusdg <harshvardhan.chandira@gmail.com>
Date: Sun, 14 Nov 2021 07:13:42 +0530
Subject: [PATCH] Added num of layers in net and netRT,suggestions from pull
 request 270 on main repo

---
 CMakeLists.txt            | 26 ++++++++++++++++++++++++++
 include/tkDNN/NetworkRT.h |  2 ++
 src/NetworkRT.cpp         |  6 +++++-
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index be3be48..7cc9e33 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -6,6 +6,32 @@ set(CMAKE_CXX_STANDARD 14)
 option(ENABLE_OPENCV_CUDA_CONTRIB "Enable OpenCV CUDA Contrib" OFF )
 
 find_package(CUDA 9.0 REQUIRED)
+if (CUDA_FOUND)
+    set(OUTPUTFILE ${CMAKE_CURRENT_SOURCE_DIR}/cmake/cuda_script) # No suffix required
+    execute_process(COMMAND "rm ${OUTPUTFILE}")
+    set(CUDAFILE ${CMAKE_CURRENT_SOURCE_DIR}/cmake/getCudaArch.cu)
+    execute_process(COMMAND ${CUDA_NVCC_EXECUTABLE} -lcuda ${CUDAFILE} -o ${OUTPUTFILE})
+    execute_process(COMMAND ${OUTPUTFILE}
+            RESULT_VARIABLE CUDA_RETURN_CODE
+            OUTPUT_VARIABLE ARCH)
+
+    if(${CUDA_RETURN_CODE} EQUAL 0)
+        set(CUDA_SUCCESS "TRUE")
+    else()
+        set(CUDA_SUCCESS "FALSE")
+    endif()
+
+    if (${CUDA_SUCCESS})
+        message(STATUS "CUDA Architecture: ${ARCH}")
+        message(STATUS "CUDA Version: ${CUDA_VERSION_STRING}")
+        message(STATUS "CUDA Path: ${CUDA_TOOLKIT_ROOT_DIR}")
+        message(STATUS "CUDA Libararies: ${CUDA_LIBRARIES}")
+        message(STATUS "CUDA Performance Primitives: ${CUDA_npp_LIBRARY}")
+        set(CUDA_NVCC_FLAGS "${ARCH}")
+    else()
+        message(WARNING ${ARCH})
+    endif()
+endif()
 
 SET(CUDA_SEPARABLE_COMPILATION ON)
 
diff --git a/include/tkDNN/NetworkRT.h b/include/tkDNN/NetworkRT.h
index cd297ab..ddd3d9f 100644
--- a/include/tkDNN/NetworkRT.h
+++ b/include/tkDNN/NetworkRT.h
@@ -24,6 +24,8 @@
 #include <pluginsRT/UpsampleRT.h>
 #include <pluginsRT/YoloRT.h>
 
+
+
 namespace tk { namespace dnn {
 
 using namespace nvinfer1;
diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp
index 709b9b2..2e9eb6e 100644
--- a/src/NetworkRT.cpp
+++ b/src/NetworkRT.cpp
@@ -183,7 +183,11 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
     output_dim.h = oDim.d[1];
     output_dim.w = oDim.d[2];
     output_dim.print();
-	
+    if(builderActive){
+        std::cout<<"NUMBER OF LAYERS IN NETWORK : "<<networkRT->getNbLayers()<<std::endl;
+    }
+    std::cout<<"NUMBER OF LAYERS IN ENGINE : "<<engineRT->getNbLayers()<<std::endl;
+
     // create GPU buffers and a stream
     for(int i=0; i<engineRT->getNbBindings(); i++) {
         Dims dim = engineRT->getBindingDimensions(i);