Hooks for functions?

I’ve seen that the hook methods are only available for classes that inherit from nn.Module. Is there any easy/simple way to use/do something like the hooks but in functions such as torch.cat or torch.nn.functional.interpolate? I’m particularly interested in tracing the output of these functions.

Not that I’m aware of, but since you will get the output directly using the functional API, you could simple store this output instead of using hooks or wouldn’t it work for your use case?

If you don’t use JIT, why not do simple monkey patching.

_cat = torch.cat
def cat(*a):
  return _cat(*a)
torch.cat = cat