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.).