torch.nn.Module.register_forward_hook
Module.register_forward_hook(hook: ForwardHook): RemovableHandleRegister a forward hook on this module. The hook will be called after forward() returns.
Parameters
hookForwardHook- Function called with (module, input, output)
Returns
RemovableHandle– Handle to remove the hook