Topk
torch_topk(self, k, dim = -1L, largest = TRUE, sorted = TRUE)
self
: (Tensor) the input tensor.k
: (int) the k in "top-k"dim
: (int, optional) the dimension to sort alonglargest
: (bool, optional) controls whether to return largest or smallest elementssorted
: (bool, optional) controls whether to return the elements in sorted orderReturns 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, where the indices
are the indices of the elements in the original input
tensor.
The boolean option sorted
if TRUE
, will make sure that the returned k
elements are themselves sorted
if (torch_is_installed()) { x = torch_arange(1., 6.) x torch_topk(x, 3) }
Useful links