Which function can do this?

A_selected, idxs = torch.min(A,dim=1) #A.shape = [b,2,h,w]

If I have ‘A’ and ‘idxs’, What should I do to get ‘A_selected’ without min function?

A_selected = torch.min(A,dim=1)

I’m sorry, there’s something wrong with my question

import torch
x = torch.rand(1, 2, 3, 4)
answer, index = torch.min(x, dim=1)

my_result = torch.gather(x, 1, index[:, None, ...]).squeeze(1)
print((my_result == answer).all())  # True
1 Like

Problem solved, Thank you soooooo much~~! :laughing: