TensorRT: ITopKLayer doesn't work like expected with tensorrt 8

Description

When I use the topk layer with a semantic segmentation network converted from pytorch to tensorrt (pidnet) using tkdnn , it works properly with tensorrt 7 , but the same code when used with tensorrt 8 produces a constant 0 output, is this a bug with tensorrt or was the topk layer updated significantly from tensorrt 7 to 8 ,requiring me to change the code for tensorrt8 ?

Environment

TensorRT Version: 7.2.3/8.5 NVIDIA GPU: RTX 3060 Laptop NVIDIA Driver Version: 525.78.01 CUDA Version: 11.1/11.4 CUDNN Version: 8.1/8.6 Operating System: Ubuntu 20.04 Python Version (if applicable): Tensorflow Version (if applicable): PyTorch Version (if applicable): Baremetal or Container (if so, version):

Relevant Files

Steps To Reproduce

To reproduce the results clone my fork of the tkdnn repo at this commit 91a4a3feb756c77480a0a473745e42216c8ec576 and build test_pidnet_s. The weights for building the network can be found here