Hooks, dynamo, and unsupported builtins

Hello!

I have some functions that I’d like to attach as pre/post hooks to instances of Module, but when I do so, I get “UserWarning: Graph break due to unsupported builtin”. I’d like to be able to use these hooks for forward_pre, forward, full_backward_pre, and full_backward. None of the code in the callbacks should be analyzed by the compiler, but the calls can’t be optimized out.

How would I go about this?

What ended up seeming to work is decorating the hooks with @torch._dynamo.disable. Is there any reason not to use this approach?

For reference, I want to be able to execute arbitrary C code in these hooks.

Marking the hooks with disable is probably reasonable (it accomplishes your goal of running the custom c-bound code in eager mode while compiling other parts of your model)

If you’re willing to do a bit more work to make your model compile friendly, you can wrap your custom c function into a PyTorch custom op as per this guide: PyTorch Custom Operators — PyTorch Tutorials 2.6.0+cu124 documentation

This will avoid graph breaks and generally improve performance (and can open the doors to eg exporting your model and running in a no-python environment if that’s something you want)