Thanks for your feedback.
Let’s just ignore the batch issue in here.
Says we have a segmentation map, which is 320x320. And the model’s prediction could be in a one-hot fashion, which is 21x320x320 (the 21 is the class category).
At the moment, I want to find the top 3 maximize elements along the first dimension (i.e., 21), for the entire map. Meanwhile, I want all the other’s non-selected elements filled to zero.  Therefore, the result remains the same architecture, as 21x320x320.
It will be more like a topk mask for the original one-hot result.