Calculating the loss for each class in Multiclass segmentation with weights

Question

Calculating the loss with pytorch is easy using the reduction feature:

criterion = nn.CrossEntropyLoss(weight=cw, reduc