torch.index_select
torch.index_select(input, dim, index, *, out=None) -> Tensor
Returns a new tensor which indexes the "input" tensor along
   dimension "dim" using the entries in "index" which is a
   LongTensor.
The returned tensor has the same number of dimensions as the
   original tensor ("input").  The "dim"th dimension has the same size
   as the length of "index"; other dimensions have the same size as in
   the original tensor.
Note:
 The returned tensor does **not** use the same storage as the
 original tensor.  If "out" has a different shape than expected,
 we silently change it to the correct shape, reallocating the
 underlying storage if necessary.

Parameters:
      * input (Tensor) -- the input tensor.
  * **dim** (*int*) -- the dimension in which we index

  * **index** (*IntTensor** or **LongTensor*) -- the 1-D tensor
    containing the indices to index

Keyword Arguments: