if you have some intermediate layer
layer(x : Tensor[b,in], p : Tensor[*]) → y : Tensor[b,out]
you don’t have to collect gradients wrt p by batch element to rescale p.grad - you can rescale output gradient (2d view in this case) rowwise, with scalar weights. But note:

loss reduction=‘none’ + weighted sum of losses approach has the same effect, applied to the whole network.
I have doubts about mixing different sample weights in one backpropagation, but I’ll assume you know what you’re doing
this reweighting sticks for earlier layers, i.e. for above function gradient wrt x is affected too. This can be counteracted with another reweighting (register_hook or similar approach).