Getting certain values of a tensor from indices

Hey guys,

I have a tensor of shape (5,10,15):
x = torch.rand(5,10,15)

and then I would like to get the indices of max values along last column:
ids = torch.argmax(x, dim=2) # this has shape (5,10)

Now, what do I do to get the values of the corresponding indices?
My output would be of shape (5,10,1)

u can make shape (5,10,1) with

ids = torch.argmax(x, dim=2).unsqueeze(dim=2)

ok, but how do I fetch those x values that match the ids?
my output of shape (5,10,1) are only the max values of x with respect to last dimension.

In a simple 2D case:
x = [ [1,2,3],
[4, 5, 2]]
ids = [2,1] # 2 because of fetching ‘3’ and 1 because of fetching ‘5’
then i want my output to be
[3,5]. Right?

x = tc.randn(5, 10, 15)
indexes = x.argmax(2)
out = tc.gather(x, 2, indexes.unsqueeze(2))

Awesome! that solves the problem, thanks G.M

Small followup tho. lets say you have
x = tc.randn(5,10,15,20)
indexes # shape of 5,10

so now the indexes tell me which guy from dim=2 is the best, but then this best guy has a dimenions of not 1, but 20. How can apply indexes to x now? In a for loop it would be like:

for i
    for j
        out = x[i,j,indexes[i,j],:]

The final shape of out is 5,10,20.

The same code works here, too.

in order to use gather i need to have x and ids of the same shape except on dim.
in my new case the shapes are
x - (5,10,15,20)
ids - (5,10)
I can use unsqueeze to make shape of ids either (5,10,1) or (5,10,1,1) but it still wont match the shape of x.

>>> import torch as tc
>>> a=tc.randn(5,10,15,20,25)
>>> b=a.argmax(2)
>>> c=tc.gather(a,2,b.unsqueeze(2))