From 1ea6097a9704ecb6b83406a2541d51b1819a76cc Mon Sep 17 00:00:00 2001
From: Francesco Gatti <gattifrancesco@hotmail.it>
Date: Thu, 23 Apr 2020 23:51:23 +0200
Subject: [PATCH] Reshape batch fix

---
 include/tkDNN/pluginsRT/ReshapeRT.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/include/tkDNN/pluginsRT/ReshapeRT.h b/include/tkDNN/pluginsRT/ReshapeRT.h
index 36c8c63..97030db 100644
--- a/include/tkDNN/pluginsRT/ReshapeRT.h
+++ b/include/tkDNN/pluginsRT/ReshapeRT.h
@@ -40,7 +40,7 @@ public:
 		dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
 		dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
 
-		checkCuda( cudaMemcpy(dstData, srcData, c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice));
+		checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
 		return 0;
 	}