If you are using CrossEntropyLoss 1, my idea is that you could change the weight parameter on the fly for your batch. Something along the lines of:
criterion = nn.CrossEntropy()

....
for input, target in dataset:
    cel_weights = calculate_weights(target)
    criterion.weights = cel_weights
    ...

If you are not using CEL, maybe check if your criterion has the weight parameter, so you could change it.