stream in TRT plugin

This commit is contained in:
Francesco Gatti
2017-08-10 19:21:15 +02:00
parent 9a6058ac4a
commit 3124f86878
9 changed files with 25 additions and 20 deletions
+2 -2
View File
@@ -17,12 +17,12 @@ void activation_leaky(dnnType *input, dnnType *output, int size) {
/**
ELU activation function
*/
void activationLEAKYForward(dnnType* srcData, dnnType* dstData, int size)
void activationLEAKYForward(dnnType* srcData, dnnType* dstData, int size, cudaStream_t stream)
{
int blocks = (size+255)/256;
int threads = 256;
activation_leaky<<<blocks, threads>>>(srcData, dstData, size);
activation_leaky<<<blocks, threads, 0, stream>>>(srcData, dstData, size);
}