All optimizers I know about require gradients wrt parameters which are always leaf nodes. So, it shouldn’t be a problem for implementing an optimizer. In my case I need those gradients only for debugging purposes.
For those are interested, I ended up doing like this:
class LazyContainer(object):
    def __call__(self, g):
        self.grad_norm = g.norm(2)
container = LazyContainer()
h.register_hook(container)
# later
print(container.grad_norm)