Can we set output dim to 1D in "getOutputDimensions" function?

Hi,

I’m writing a program that converts a trained TF network to Tensorrt. I have met so many issues.

The raw input is a set of points with shape [batch_size, feature_channel=3 (e.g.: x,y,z), points_size]

In one of the custom ops(plugin), the output shape of the operation should be a list of indices with shape: [batch_size, indices_size]

I’m confused about how can we set the output dimension to 1D in the “getOutputDimensions” function.

Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;

Here, it should returns

Dims2(batch_size, indices_size)

or

Dims(indices_size) //Error

.

In the code of example plugins, it seems that the batch_size should not be included in the returned dims. But I also found that the returned dims must have at least two dimensions. How can I resolve it? Thanks in advance!

After some search, I found a way to set Dims to 1D.

Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
{
    ASSERT(index == 0);
    ASSERT(nbInputDims == 1);

    Dims output;
    output.nbDims = 1;
    output.d[0] = indices_size;
    return output;
}

Am I correct? Can anyone confirm this method could work?

Could you please let us know if you are still facing this issue?

Thanks