I’m not sure someone asked it before…
Says we have a one-hot segmentation mask with BxNxHxW, where B=batch, N=class-categories, H=height, W=width of the mask. I’m trying to find the top 3 maximize results from each one-hot element-wise vector of the map, and let other unselected elements be 0 (which means we want the dimension not to change).
It is like a keepdim version of the torch.topk.
I didn’t find a way to solve this…   I appreciate any advice in here, thanks.