Hi,
I am afraid these functions are at a much lower level than torch.nn and you won’t be able to use a direct hooking system.
One thing you can try to do is wrap your input in a Tensor-like object and write a custom implementation that calls your hook and then delegate to the original pytorch implementation: Extending PyTorch — PyTorch 1.8.1 documentation 1