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).