Solved by ptrblck in post #2 
              
              
                You could avoid the reduction in your criterion by using: 
criterion = nn.CrossEntropyLoss(reduction='none')

This will give you a loss values for each sample in the batch. 
After weighting each element, you could take the average (or normalize, sum etc.).