Solved by ayalaa2 in post #2 
              
              
                I would set the reduction to none: CrossEntropyLoss(reduction='none') 
This returns you a num_batch size tensor. So you can then do something like this: 
(weights * criterion(output, labels)).mean()