TensorRT format handling

I am trying to build an SSD network with TensorRT’s C++ API. I am wondering about how TensorRT handles formats. The last part of SSD involves concatenation of a bunch of tensors. In my understanding, TensorRT is free to pick the intermediate format of all these tensors. Then, how should I supply the axis argument to IConcatenationLayer, since I don’t know what the layout is going to be? is there a way to force an intermediate layer to have a certain format, similar to how you can do network->getOutput(i)->setAllowedFormats(formats); ?

TensorRT uses NCHW uniformly when defining the semantics of operations. You can use the TensorFormat enum to gain access to TensorRT’s internal data layouts at network boundaries, which are optimized for TensorCore. None of these formats is strictly NHWC.

Please refer to below links in case it helps:

Thanks