Yes, there are some changes since TAO 4.0.
Currently for UNet,
GitHub - NVIDIA-AI-IOT/deepstream_tao_apps: Sample apps to demonstrate how to deploy models trained with TAO on DeepStream,
argmax_1/output : A [batchSize, H, W, 1] tensor containing the class id per pixel location
Previously for Unet,
GitHub - NVIDIA-AI-IOT/deepstream_tao_apps at release/tao3.0_ds6.2ga,
softmax_1: A [batchSize, H, W, C] tensor containing the scores for each class
Related topic:
DeepStream compatibility issues with UNET output layer change - #28 by adityasingh.