From 714d0685a911d5cc483e3797e603154e6848d9dc Mon Sep 17 00:00:00 2001
From: oren <virtualoren@gmail.com>
Date: Sat, 5 Dec 2020 16:52:47 -0600
Subject: [PATCH] Initiliaze dataDim_t with nvinfer1::Dims

---
 include/tkDNN/Network.h |  5 +++++
 src/NetworkRT.cpp       | 14 ++++----------
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/include/tkDNN/Network.h b/include/tkDNN/Network.h
index b78acff..571e728 100644
--- a/include/tkDNN/Network.h
+++ b/include/tkDNN/Network.h
@@ -3,6 +3,7 @@
 
 #include <string>
 #include "utils.h"
+#include "NvInfer.h"
 
 namespace tk { namespace dnn {
 
@@ -20,6 +21,10 @@ struct dataDim_t {
 
     dataDim_t() : n(1), c(1), h(1), w(1), l(1) {};
 
+    dataDim_t(nvinfer1::Dims &d) :
+        n(1), c(d.d[0] ? d.d[0] : 1), h(d.d[1] ? d.d[1] : 1),
+        w(d.d[2] ? d.d[2] : 1), l(d.d[3] ? d.d[3] : 1) {};
+
     dataDim_t(int _n, int _c, int _h, int _w, int _l = 1) :
         n(_n), c(_c), h(_h), w(_w), l(_l) {};
 
diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp
index 9d38440..889fb9f 100644
--- a/src/NetworkRT.cpp
+++ b/src/NetworkRT.cpp
@@ -167,25 +167,19 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
 
 
     Dims iDim = engineRT->getBindingDimensions(buf_input_idx);
-    input_dim.n = 1;
-    input_dim.c = iDim.d[0];
-    input_dim.h = iDim.d[1];
-    input_dim.w = iDim.d[2];
+    input_dim = dataDim_t(iDim);
     input_dim.print();
 
     Dims oDim = engineRT->getBindingDimensions(buf_output_idx);
-    output_dim.n = 1;
-    output_dim.c = oDim.d[0];
-    output_dim.h = oDim.d[1];
-    output_dim.w = oDim.d[2];
+    output_dim = dataDim_t(oDim);
     output_dim.print();
 	
     // create GPU buffers and a stream
     for(int i=0; i<engineRT->getNbBindings(); i++) {
         Dims dim = engineRT->getBindingDimensions(i);
-        buffersDIM[i] = dataDim_t(1, dim.d[0], dim.d[1], dim.d[2]);
+        buffersDIM[i] = dataDim_t(dim);
         std::cout<<"RtBuffer "<<i<<"   dim: "; buffersDIM[i].print();
-        checkCuda(cudaMalloc(&buffersRT[i], engineRT->getMaxBatchSize()*dim.d[0]*dim.d[1]*dim.d[2]*sizeof(dnnType)));
+        checkCuda(cudaMalloc(&buffersRT[i], engineRT->getMaxBatchSize()*buffersDIM[i].tot()*sizeof(dnnType)));
     }
     checkCuda(cudaMalloc(&output, engineRT->getMaxBatchSize()*output_dim.tot()*sizeof(dnnType)));
 	checkCuda(cudaStreamCreate(&stream));