savourylie:

topk, indices = torch.topk(activation, 2, dim=1)


I was able to reduce the complexity to one single for loop using a selection tensor:
activation = x.sum(dim=2).sum(dim=2)
topk, indices = torch.topk(activation, self.k, dim=1)
selection_tensor = torch.zeros_like(x)
            
for i, _ in enumerate(indices):
    selection_tensor[i, indices[i], :, :] = 1
            
x = x * selection_tensor

This helps a lot but still it’d be nice if we can throw away the loop all together.