Thanks for your quick reply!



 ptrblck:

I tried to stick to your approach and as you can see, both methods yield the same gradients.


Yes, with register_hood. Both method yield the same gradients. Before this, I have tried:



 ptrblck:

print(torch.allclose(modelA.model.features[0].weight.grad), modelB.features[0].weight.grad))


I think the gradients of ModelOutputs(net, layer=40) is different from net.features[40].weight.grad.