How to register forward hooks for each module

2

You could use model.named_modules() to iterate all submodules including their names and store the hooks in a dict:

self.hooks = {}
for name, module in model.named_modules():
    hooks[name] = module.register_forward_hook(self,hook_fn)
7 Likes