In the batch gradient method, the parameters are updated by simply averaging the gradients in the batch direction, but in my particular situation, I need to change the aggregation function (average → weighted sum, etc.).
How can I make this happen? As far as I can tell, it’s not possible with register_full_backward_hook.