torch.nn.modules.module.register_module_forward_hook
torch.nn.modules.module.register_module_forward_hook(hook)
Registers a global forward hook for all the modules
Warning:
 This adds global state to the *nn.module* module and it is only
 intended for debugging/profiling purposes.

The hook will be called every time after "forward()" has computed
   an output. It should have the following signature:
  hook(module, input, output) -> None or modified output

The input contains only the positional arguments given to the
   module. Keyword arguments won't be passed to the hooks and only to
   the "forward". The hook can modify the output. It can modify the
   input inplace but it will not have effect on forward since this is
   called after "forward()" is called.
Returns:
      a handle that can be used to remove the added hook by calling
      "handle.remove()"
Return type:
      "torch.utils.hooks.RemovableHandle"