Best loss function for special classification problem

Question

I build a Conv1D model in pytorch to classify items into 6 categories (from 0 to 5).

I use therefore nn.CrossEntropyLoss that penalizes the errors with high probabilities