Yes, it is possible. Refer to SegFormer - NVIDIA Docs and also download latest notebook from TAO Toolkit Getting Started | NVIDIA NGC. More info can be also found in TAO 4 Segformer Input and output dimensions and tensors - #12 by Morganh.
Will sync internally for this feature request. Thanks.