diff --git a/CMakeLists.txt b/CMakeLists.txt
index ce0f405..19ab58e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -11,11 +11,12 @@ cuda_add_library(kernels SHARED src/kernels/activation_elu.cu
                                 src/kernels/reorg.cu
                                 src/kernels/softmax.cu)
 
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -std=c++11")
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include ${CUDA_INCLUDE_DIRS})
 add_library(tkDNN SHARED src/Layer.cpp src/LayerWgs.cpp
     src/Dense.cpp src/Activation.cpp src/Conv2d.cpp src/Flatten.cpp src/MulAdd.cpp src/Pooling.cpp src/Softmax.cpp
     src/Route.cpp src/Reorg.cpp src/Region.cpp src/Network.cpp src/utils.cpp)
-target_link_libraries(tkDNN kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} -lcudnn)
+target_link_libraries(tkDNN kernels ${CUDA_LIBRARIES} ${CUDA_CUBLAS_LIBRARIES} -lcudnn -lnvinfer)
 
 add_executable(test_simple tests/test/test.cpp)
 target_link_libraries(test_simple tkDNN)
diff --git a/include/kernels.h b/include/kernels.h
index 3901099..97f81c6 100644
--- a/include/kernels.h
+++ b/include/kernels.h
@@ -1,10 +1,10 @@
 #include "utils.h"
-#include "Layer.h"
 
 void activationELUForward(value_type* srcData, value_type* dstData, int size);
 void activationLEAKYForward(value_type* srcData, value_type* dstData, int size);
 void activationLOGISTICForward(value_type* srcData, value_type* dstData, int size);
 
-void reorgForward(value_type* srcData, value_type* dstData, tkDNN::dataDim_t dim, int stride);
+void reorgForward(  value_type* srcData, value_type* dstData, 
+                    int n, int c, int h, int w, int stride);
 void softmaxForward(float *input, int n, int batch, int batch_offset, 
                     int groups, int group_offset, int stride, float temp, float *output);
diff --git a/src/Conv2d.cpp b/src/Conv2d.cpp
index 0ce544a..46360d6 100644
--- a/src/Conv2d.cpp
+++ b/src/Conv2d.cpp
@@ -37,7 +37,7 @@ Conv2d::Conv2d( Network *net, dataDim_t in_dim, int out_ch,
                 paddingH, paddingW, // padding
                 strideH, strideW, // stride
                 1,1, // upscale
-                CUDNN_CROSS_CORRELATION) );
+                CUDNN_CROSS_CORRELATION, CUDNN_DATA_FLOAT) );
 
     // find dimension of convolution output
     checkCUDNN( cudnnGetConvolution2dForwardOutputDim(
diff --git a/src/Network.cpp b/src/Network.cpp
index 6e705e6..c0c101d 100644
--- a/src/Network.cpp
+++ b/src/Network.cpp
@@ -1,4 +1,5 @@
 #include <iostream>
+#include "NvInfer.h"
 
 #include "tkdnn.h"
 #include "Network.h"
@@ -10,8 +11,10 @@ Network::Network() {
 
     float tk_ver = float(tkDNN::getVersion())/1000;
     float cu_ver = float(cudnnGetVersion())/1000;
+    float rt_ver = float(NV_TENSORRT_MAJOR) + float(NV_TENSORRT_MINOR)/10 + float(NV_TENSORRT_PATCH)/100;
 
-    std::cout<<"New NETWORK (tkDNN v"<<tk_ver<<", CUDNN v"<<cu_ver<<")\n";
+    std::cout<<"New NETWORK (tkDNN v"<<tk_ver
+             <<", CUDNN v"<<cu_ver<<", TensorRT v"<<rt_ver<<")\n";
     dataType = CUDNN_DATA_FLOAT;
     tensorFormat = CUDNN_TENSOR_NCHW;
 
@@ -32,7 +35,7 @@ value_type* Network::infer(dataDim_t &dim, value_type* data) {
     //do infer for every layer
     for(int i=0; i<num_layers; i++)
         data = layers[i]->infer(dim, data);
-        
+
     return data;
 }
 
diff --git a/src/Reorg.cpp b/src/Reorg.cpp
index 543457e..7c64f33 100644
--- a/src/Reorg.cpp
+++ b/src/Reorg.cpp
@@ -26,7 +26,7 @@ Reorg::~Reorg() {
 
 value_type* Reorg::infer(dataDim_t &dim, value_type* srcData) {
 
-    reorgForward(srcData, dstData, dim, stride);
+    reorgForward(srcData, dstData, dim.n, dim.c, dim.h, dim.w, stride);
 
     dim = output_dim;
     return dstData;
diff --git a/src/kernels/reorg.cu b/src/kernels/reorg.cu
index 37abc24..2b7cca0 100644
--- a/src/kernels/reorg.cu
+++ b/src/kernels/reorg.cu
@@ -35,14 +35,15 @@ __global__ void reorg_kernel(int N, float *x, int w, int h, int c, int batch, in
 /**
     reorg function function
 */
-void reorgForward(value_type* srcData, value_type* dstData, tkDNN::dataDim_t dim, int stride)
-{
-    int size = dim.tot();
+void reorgForward(value_type* srcData, value_type* dstData, 
+                  int n, int c, int h, int w, int stride) {
+
+    int size = n*c*h*w;
 
     int blocks = (size+255)/256;
     int threads = 256;
     
-    reorg_kernel<<<blocks, threads>>>(size, srcData, dim.w, dim.h, dim.c, dim.n, stride, false, dstData);
+    reorg_kernel<<<blocks, threads>>>(size, srcData, w, h, c, n, stride, false, dstData);
     checkCuda( cudaDeviceSynchronize() );
 }