logits is a (batch, classes) tensor of raw multi-class scores. Return the predicted class index for each sample in the batch (softmax then argmax along the class dimension).

Input

tensor([[2.0000, 1.0000, 0.1000],
        [0.1000, 3.0000, 0.2000],
        [1.0000, 1.0000, 5.0000]])

Output

tensor([0, 1, 2])