CPU output from custom plugin

I want to pass an CPU data-dependent variable from one custom layer to another. Is it possible?
I can upload a variable to GPU 1x1 tensor and download it to CPU in another layer, but it is non-elegant in my opinion…

Hi @arsenium.ml ,
Can you please elaborate on the issue.


I want to return tensor with data-dependent shape from TensorRT plugin. Suppose this tensor have shape [rows_real, 3], and rows_real is data-dependent and can be determined in enqueue(…) function (not in getOutputDimensions).

As I can understand, the only option now is to set output tensor shape in getOutputDimensions(…) to [rows_upper_bound, 3]. Next, in enqueue(…) I calculate rows_real (it is CPU variable now), and populate the first rows_real rows of output tensor via cudaMemcpy(…).

But I still need to pass rows_real (CPU variable) to the next layer in order it can slice it input properly. How do I can do this properly?

One option is to make the second output from plugin (call is output_tensor2) with the shape 1x1 and populate is with rows_real (again via cudaMemcpy). In the next layer (TensorRT plugin), I copy input_tensor2[0] element to CPU, name it rows_real and do the slice of input_tensor1.

So the question is - is there an option to pass CPU variable rows_real from layer1 to layer2 as CPU tensor (without copying to and from GPU)?