How can a custom plugin layer change the batch size for its later part (e.g. ROI pooling) ?

For example, in FasterRCNN, the “ROI pooling” layer takes two inputs, a 1xCxHxW feature map and 1x300x4 ROIs, where batch size is N=1 (a single image) and 300 ROI box coordinates are given.

The output of ROI pooling will be 300x512x7x7, where there are 300 “pooled” features of 512x7x7. The rest of the network (that is, region classification part) will take this tensor and should regard 300 as a batch size.

Let’s say I implement the “ROI pooling” custom plugin layer. I tried to set its output dimension as below but it didn’t work. It seems that all Plugin layers assume the first (batch) dimension is hidden, and act like only DimsCHW is valid. Is this true? Is there a way to set a new batch size?

class CropPooling : public IPlugin
    int getNbOutputs() const override { return 1; }

    Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) {
        assert(index == 0);
        return DimsNCHW(300, 512, 7, 7);

I checked “sampleFasterRcnn” but unfortunately there was no source code. I also took a look at example, but all the custom layers were dealing with just DimsCHW.

Please advice/help.