diff --git a/.gitignore b/.gitignore
index 446d490..02f2a8e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,3 +8,4 @@ build/
 *.h5
 *.tar.gz
 *.weights
+.idea/
\ No newline at end of file
diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp
index cfc1fa3..a1ed830 100644
--- a/src/NetworkRT.cpp
+++ b/src/NetworkRT.cpp
@@ -183,6 +183,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {
     if(type == LAYER_UPSAMPLE)
         return convert_layer(input, (Upsample*) l);
 
+    std::cout<<l->getLayerName()<<"\n";
     FatalError("Layer not implemented in tensorRT");
     return NULL;
 }
diff --git a/tests/simple/test_simple.cpp b/tests/simple/test_simple.cpp
index a7628ea..d331b96 100644
--- a/tests/simple/test_simple.cpp
+++ b/tests/simple/test_simple.cpp
@@ -16,7 +16,6 @@ int main() {
     tk::dnn::Activation l1(&net, CUDNN_ACTIVATION_RELU);
     tk::dnn::Conv2d     l2(&net, 4, 2, 2, 1, 1, 0, 0, c1_bin);
     tk::dnn::Activation l3(&net, CUDNN_ACTIVATION_RELU);
-    tk::dnn::Flatten    l4(&net);
     tk::dnn::Dense      l5(&net, 4, d2_bin);
     tk::dnn::Activation l6(&net, CUDNN_ACTIVATION_RELU);
 
@@ -25,23 +24,45 @@ int main() {
     dnnType *input_h;
     readBinaryFile(input_bin, dim.tot(), &input_h, &data);
 
+    // Print input
+    std::cout<<"\n======= INPUT =======\n";
     printDeviceVector(dim.tot(), data);
-    dim.print(); //print initial dimension
-    
-    TIMER_START
-    // Inference
-    data = net.infer(dim, data); dim.print();
-    TIMER_STOP
+    std::cout<<"\n";
 
-    // Print result
-    std::cout<<"\n======= RESULT =======\n";
-    printDeviceVector(dim.tot(), data);
+    //convert network to tensorRT
+    tk::dnn::NetworkRT netRT(&net, "simple.rt");
 
-    // Print real test
-    std::cout<<"\n==== CHECK RESULT ====\n";
-    dnnType *out;
-    dnnType *out_h;
-    readBinaryFile(output_bin, dim.tot(), &out_h, &out);
-    printDeviceVector(dim.tot(), out);
+    dnnType *out_data, *out_data2; // cudnn output, tensorRT output
+
+    tk::dnn::dataDim_t dim1 = dim; //input dim
+    printCenteredTitle(" CUDNN inference ", '=', 30); {
+        dim1.print();
+        TIMER_START
+        out_data = net.infer(dim1, data);
+        TIMER_STOP
+        dim1.print();
+    }
+
+    tk::dnn::dataDim_t dim2 = dim;
+    printCenteredTitle(" TENSORRT inference ", '=', 30); {
+        dim2.print();
+        TIMER_START
+        out_data2 = netRT.infer(dim2, data);
+        TIMER_STOP
+        dim2.print();
+    }
+
+    std::cout<<"\n======= CUDNN =======\n";
+    printDeviceVector(dim.tot(), out_data);
+    std::cout<<"\n======= TENSORRT =======\n";
+    printDeviceVector(dim.tot(), out_data2);
+
+    printCenteredTitle(" CHECK RESULTS ", '=', 30);
+    dnnType *out, *out_h;
+    int out_dim = net.getOutputDim().tot();
+    //readBinaryFile(output_bin, out_dim, &out_h, &out);
+    //std::cout<<"CUDNN vs correct"; checkResult(out_dim, out_data, out);
+    //std::cout<<"TRT   vs correct"; checkResult(out_dim, out_data2, out);
+    std::cout<<"CUDNN vs TRT    "; checkResult(out_dim, out_data, out_data2);
     return 0;
 }