Hi. You can probably do with reduction ‘none’. Something like this should probably work.
criterion = CrossEntropyLoss(reduction='none')
loss = criterion(network_out, label) --> shape batch_size x 1
loss = loss * batch_weights
loss = torcn.mean(loss)