From c95c2dbfbc12357ad26b041637a295a94d861c2e Mon Sep 17 00:00:00 2001
From: Francesco Gatti <gattifrancesco@hotmail.it>
Date: Wed, 9 Aug 2023 02:28:25 +0800
Subject: [PATCH] rtinference nstep

---
 tests/test_rtinference/rtinference.cpp | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/test_rtinference/rtinference.cpp b/tests/test_rtinference/rtinference.cpp
index bcdc739..e2f4457 100644
--- a/tests/test_rtinference/rtinference.cpp
+++ b/tests/test_rtinference/rtinference.cpp
@@ -13,6 +13,10 @@ int main(int argc, char *argv[]) {
     if(argc >2)
         BATCH_SIZE = atoi(argv[2]);
 
+    int NTEST = 100;
+    if(argc >3)
+        NTEST = atoi(argv[3]);
+
     //always same test 
     srand (0); 
 
@@ -35,7 +39,7 @@ int main(int argc, char *argv[]) {
     std::vector<double> stats;
     printCenteredTitle(" TENSORRT inference ", '=', 30); 
     float total_time = 0;
-    for(int i=0; i<64; i++) {
+    for(int i=0; (NTEST > 0 ? i<NTEST : true) ; i++) {
 
         // generate input
         for(int j=0; j<netRT.input_dim.tot(); j++) {