Hi,
I would see two ways:

Return the output of each layer then use autograd.grad to backward only the part of the network that you want.
Use register_hook() on the output of each layer and have the hook function do the calculation you want and return the new gradient value from the hook.