background
Use jit.trace to export the torchscript model, and then use tvm import for compilation and acceleration. However, sometimes the inference results of tvm and pytorch are different, and it is difficult to check which layer the difference comes from. Therefore, it is currently necessary to obtain the output of the intermediate node of the torchscript model to check which layer the difference comes from.
question
I learned that register_forward_hook can implement such a function, but it does not work with RecursiveScriptModule. I am very new to pytorch, I would like to ask how much work is required to make RecursiveScriptModule support functions like register_forward_hook, and what are the approximate steps. Also since this is the debug version, I don’t care about the performance of inference.