torch.take_along_dim
torch.take_along_dim(input, indices, dim, *, out=None) -> Tensor
Selects values from "input" at the 1-dimensional indices from
   "indices" along the given "dim".
Functions that return indices along a dimension, like
   "torch.argmax()" and "torch.argsort()", are designed to work with
   this function. See the examples below.
Note:
 This function is similar to NumPy's *take_along_axis*. See also
 "torch.gather()".

Parameters:
      * input (Tensor) -- the input tensor.
  * **indices** (*tensor*) -- the indices into "input". Must have
    long dtype.

  * **dim** (*int*) -- dimension to select along.

Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> t = torch.tensor([[10, 30, 20], [60, 40, 50]])
  >>> max_idx = torch.argmax(t)
  >>> torch.take_along_dim(t, max_idx)
  tensor([60])
  >>> sorted_idx = torch.argsort(t, dim=1)
