when some Function needs to keep an input tensor, it is stored locally in saved_tensors - this is the same for built-in operators, they just use c++ version of Function. There is an obvious disconnect between Module and Function, plus some gradients are computed without storing function inputs.
For module-level hooks, it would be way easier to exploit python’s dynamism and additionally hook Module.__call__ to record inputs. In fact, I think Module.register_forward_pre_hook  exists for this usecase (again, I haven’t used it in practice).