Valid approach? LogSoftmax during training, Softmax during inference

I am training a classifier assigning one of four possible classes to each frame in a preprocessed audio stream using pytorch. I am using cross-entropy loss as the loss function for training. It is implemented as a combination of LogSoftmax and negative log-likelihood loss. During inference however I am using a softmax layer before the output, omitting the logarithm. Now I am wondering whether this yields any problems I might overlook. I suppose it should be fine, because in the end I’m deciding to choose one class and as both softmax and log-softmax are monotonic, the highest probability should belong to the same class using either. The reason I’m not simply switching to log-softmax during inference as I would loose the desirable effect of the output being an actual probability distribution if I’m not mistaken.

Are my assumptions correct so I may carry on or did I miss something important?