Shortcuts

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 return the elements in sorted order

Keyword Arguments:

out (tuple, optional) – the output tuple of (Tensor, LongTensor) that can be optionally given to be used as output buffers

Example:

>>> x = torch.arange(1., 6.)
>>> x
tensor([ 1.,  2.,  3.,  4.,  5.])
>>> torch.topk(x, 3)
torch.return_types.topk(values=tensor([5., 4., 3.]), indices=tensor([4, 3, 2]))

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources