how to get max value by the indices return by torch.max?
The interface torch.max will return value and indices， how can i use the indices to get the according elements  from another tensor?
for example:
a = torch.rand(2,3,4)
b = torch.rand(2,3,4)
# indices shape is [2, 4]
indices = torch.max(a, 1)[1]
# how to get elements by indices ?
b_max = ????


keepdim=True when calling torch.max() and torch.take_along_dim() should do the trick.
&gt;&gt;&gt; import torch
&gt;&gt;&gt; a=torch.rand(2,3,4)
&gt;&gt;&gt; b=torch.rand(2,3,4)
&gt;&gt;&gt; indices=torch.max(a,1,keepdim=True)[1]
&gt;&gt;&gt; b_max = torch.take_along_dim(b,indices,dim=1)

2D example:
&gt;&gt;&gt; a=torch.rand(2,3)
&gt;&gt;&gt; a
tensor([[0.0163, 0.0711, 0.5564],
        [0.4507, 0.8675, 0.5974]])
&gt;&gt;&gt; b=torch.rand(2,3)
&gt;&gt;&gt; b
tensor([[0.7542, 0.1793, 0.5399],
        [0.2292, 0.5329, 0.2084]])
&gt;&gt;&gt; indices=torch.max(a,1,keepdim=True)[1]
&gt;&gt;&gt; torch.take_along_dim(b,indices,dim=1)
tensor([[0.5399],
        [0.5329]])

