torch.topk
torch.topk(input, k, dim=None, largest=True, sorted=True, *, out=None)
Returns the "k" largest elements of the given "input" tensor along
   a given dimension.
If "dim" is not given, the last dimension of the input is chosen.
If "largest" is "False" then the k smallest elements are
   returned.
A namedtuple of (values, indices) is returned with the values
   and indices of the largest k elements of each row of the
   input tensor in the given dimension dim.
The boolean option "sorted" if "True", will make sure that the
   returned k elements are themselves sorted
Parameters:
      * input (Tensor) -- the input tensor.
  * **k** (*int*) -- the k in "top-k"

  * **dim** (*int**, **optional*) -- the dimension to sort along

  * **largest** (*bool**, **optional*) -- controls whether to
    return largest or smallest elements

  * **sorted** (*bool**, **optional*) -- controls whether to
