How to get indices of top-K values from a numpy array
Let suppose I have probabilities from a Pytorch or Keras predictions and result is with the softmax function
from scipy.special import softmax
probs = softmax(np.random.randn(20,10),1) # 20 instances and 10 class probabilities
probs

I want to find top-5 indices from this numpy array. All I want to do is to run a loop on the results something like:
for index in top_5_indices:
    if index in result:
        print('Found')

I'll get if my results are in top-5 results.
Pytorch has top-k function and I have seen numpy.argpartition but I have no idea how to get this done?

A little more expensive, but argsort would do:
idx = np.argsort(probs, axis=1)[:,-5:]

If we are talking about pytorch:
probs = torch.from_numpy(softmax(np.random.randn(20,10),1))

values, idx = torch.topk(probs, k=5, axis=-1)

