How can I force torch.jit.trace to compule my module by ignoring hooks?
I have a module containing hook, and I would like to compile it with jit's trace:

compiled_model = torch.jit.trace(model,  torch.rand(1, 3, 256, 256))


But I get the error:

ValueError: Modules that have hooks assigned can't be compiled


How can I force trace to ignore the hooks ?

If you want to bypass trace's check, you can recursively remove all the hooks from your model.

This can be done by iterating over the children:

from collections import OrderedDict
def remove_hooks(model):
    model._backward_hooks = OrderedDict()
    model._forward_hooks = OrderedDict()
    model._forward_pre_hooks = OrderedDict()
    for child in model.children():
        remove_hooks(child)


Then you can force the compilation:

remove_hooks(model)
compiled_model = torch.jit.trace(model,  torch.rand(1, 3, 256, 256))


But if the hook are actually doing real work and you want to keep them in the trace (which was my case) you can just comment torch's raise in torch/jit/__init__.py the lines :

if orig._backward_hooks or orig._forward_hooks or orig._forward_pre_hooks:
    raise ValueError("Modules that have hooks assigned can't be compiled")


It worked for me and I managed to compile a fastai model.
