Solved it with the help of this method How to select index over two dimension? and grouping my features across the sets (group1 - all row 1s among the k sets, group2 - all row 2s among the k sets). Then I apply the norm and argmax on these groups to achieve the desired output.
1 Like