Premium problem106. Top-K Predictions

Easy Locked

logits is (B, C). Return the tuple (indices, probabilities), each (B, k), giving each row's k most likely classes ordered from most to least likely, with the probabilities taken from a softmax over the classes.

This is what a classifier actually returns to a caller: the labels and how confident it is, not the raw scores.

Input

logits =
tensor([[1.0000, 3.0000, 2.0000],
        [0.5000, 0.1000, 0.9000]])
k = 2

Output

(tensor([[1, 2],
        [2, 0]]), tensor([[0.6652, 0.2447],
        [0.4718, 0.3162]]))

Premium problem

This one's part of Premium. Unlock the full PyTorch track plus every other premium problem on the site.

Implement solve(...)