Get 2d indices for max and min in pytorch

Question

lets say I have a tensor of shape [batchsize , 7 , 7] I wish to get for each sample in the batch the maximum two values with 2d indices + the minimum two values with 2d indices