Hi,
The problem here is that most Functions are actually implemented in cpp and called automatically. There is not way to get some code executed everytime a Function is forwarded (or the backward created).
In your case, you can also say that you only support nn.Module based models. This ways, you can look into the model that is given to you and use .register_forward_hook to get both the inputs and outputs of this Module.
From there, you can .register_hook on the tensors you want to get their value and gradients.
Doing that for all relevant modules like Conv and Linear will give you something similar to what you want no?