torch.gather
torch.gather(input, dim, index, *, sparse_grad=False, out=None) -> Tensor
Gathers values along an axis specified by dim.
For a 3-D tensor the output is specified by:
  out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
  out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
  out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

"input" and "index" must have the same number of dimensions. It is
   also required that "index.size(d) <= input.size(d)" for all
   dimensions "d != dim".  "out" will have the same shape as "index".
   Note that "input" and "index" do not broadcast against each other.
Parameters:
      * input (Tensor) -- the source tensor
  * **dim** (*int*) -- the axis along which to index

  * **index** (*LongTensor*) -- the indices of elements to gather

Keyword Arguments:
      * sparse_grad (bool, optional) -- If "True", gradient
        w.r.t. "input" will be a sparse tensor.