From 07193fc343062aedaf7ec5b168ab49417e4ee5a5 Mon Sep 17 00:00:00 2001
From: Francesco Gatti <gattifrancesco@hotmail.it>
Date: Tue, 16 Jun 2020 12:44:12 +0200
Subject: [PATCH] dealloc in test.h fix #36

---
 include/tkDNN/test.h | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/include/tkDNN/test.h b/include/tkDNN/test.h
index be3d891..4e238ff 100644
--- a/include/tkDNN/test.h
+++ b/include/tkDNN/test.h
@@ -1,7 +1,7 @@
 
 #include <tkdnn.h>
 int testInference(std::vector<std::string> input_bins, std::vector<std::string> output_bins, 
-  tk::dnn::Network *net, tk::dnn::NetworkRT *netRT = nullptr) {
+    tk::dnn::Network *net, tk::dnn::NetworkRT *netRT = nullptr) {
 
     std::vector<tk::dnn::Layer*> outputs;
     for(int i=0; i<net->num_layers; i++) {
@@ -67,7 +67,11 @@ int testInference(std::vector<std::string> input_bins, std::vector<std::string>
             std::cout<<"CUDNN vs TRT    "; 
             ret_cudnn_tensorrt |= checkResult(odim, cudnn_out[i], rt_out[i]) == 0 ? 0 : ERROR_CUDNNvsTENSORRT;
         }
-    }
-    return ret_cudnn | ret_tensorrt | ret_cudnn_tensorrt;
 
-  }
\ No newline at end of file
+        delete [] out_h;
+        checkCuda( cudaFree(out) );
+    }
+    delete [] input_h;
+    checkCuda( cudaFree(data) );
+    return ret_cudnn | ret_tensorrt | ret_cudnn_tensorrt;
+}
\ No newline at end of file