is it possible to pass your own weights into the loss function during each batch such that it affects how the loss is calculated? Note: I don’t want to use the same weights for every batch (which is the weight argument), I want to weigh the loss from each output neuron dynamically based on a rule from the ground truth labels.
So for example, a 2 neuron final layer can have loss weighing [[1,0.5],[0,0.5]] for the first batch of 2, and the second batch of 2 can have [[0,0],[1,1]]… and so on.
and the code might look something like this in the training loop:
target = torch.randn([2, 2], dtype=torch.float32) #placeholder, will be taken from data loader
weights = create_weights_from_target(target)
output = model(inputs)
criterion = torch.nn.BCEWithLogitsLoss()
criterion(output, target, weights=weights) #extra arg to specify weights for each batch

Short of reimplementing the whole loss function, is there a better way to do this?
Please let me know if you need me to clarify anything. Thank you very much!!!