Hi @Yaroslav_Bulatov!
Thank you for the very quick and detailed response!
I assume the reason behind why they’re equivalent is because the gradient of mean(L_i) w.r.t parameter_j  is equal to the gradient of mean(L_i) w.r.t L_i’ * gradient L_i’ w.r.t parameter_j and the first gradient term is just 1, so they’re equivalent?  (I assume this only works for independent samples?)
Thank you for the explanation between per-sample gradient and “gradient”, I think I fully understand it now!
The follow-up question I would have, if that’s ok, is when I did an example script for my problem I got grad1.mean(dim=0) to equal param.grad. Which would show my own method is correct, right?
However, if I iterate through all my data sequentially and calculate an individual loss, and gradients via loss_i.backward() and store model.parameters.grad() within a list (with also calling model.zero_grad() between samples so grads don’t accumulate) I don’t get the same values as those returned via the hooks. So, I’m somewhat confused as to why they’re different.
Are there any caveats to using hooks for this? To give a bit more detail about my loss function it does contain a Laplacian operator, so calculates the Trace of the Hessian matrix of the output of my network w.r.t the inputs. Could this be a potential issue? Perhaps maybe the grad_output term is from the Laplacian rather than the loss function itself? Thereby explaining any potential difference between the sequential method and the batch method?
If it’s not too much could I post an example to illustrate my particular problem?
Thank you for your time!