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 int testInference(std::vector input_bins, std::vector output_bins, - tk::dnn::Network *net, tk::dnn::NetworkRT *netRT = nullptr) { + tk::dnn::Network *net, tk::dnn::NetworkRT *netRT = nullptr) { std::vector outputs; for(int i=0; inum_layers; i++) { @@ -67,7 +67,11 @@ int testInference(std::vector input_bins, std::vector 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