I believe the first one is much better. The squashing function does not change the results of inference; i.e., if you pick the class with the highest probability vs picking the class with the highest logit, you’ll get the same results.
So making a prediction is always the same, a variant of output.max(dim=1)[1]. In particular, in your first example, there is no need to compute probabilities first (unless you want to show them to the user / you are working on some edge case)