I have a (n, n) mask and corresponding (3, n, n) image. I’m trying to use the torch.topk() function to get the top k sigmoid pixel values and erase them from the corresponding image. However, I’m struggling quite a bit trying to understand how topk() deals with dimensions… To properly compute this, do I need to unsqueeze the matrix? Thanks!