How to implent tf.nn.in_top_k in pytorch

Question

I want to implent tf.nn.in_top_k in pytorch. Here is the link of tf.nn.in_top_k,

tf.math.in_top_k(
    targets, predictions, k, name=None
)

It computed precision at k as a bool Tensor and will return a Tensor of type bool.

tf.nn.in_top_k

I wonder whether there are similar api in pytorch?

Answer 1

AFAIK there is no equivalent in_top_k function built into pytorch. It's relatively straightforward to write one. For example

def in_top_k(targets, preds, k):
    topk = preds.topk(k)[1]
    return (targets.unsqueeze(1) == topk).any(dim=1)