oneflow.nn.Module.register_full_backward_hook¶
-
Module.
register_full_backward_hook
(hook: Callable[[oneflow.nn.modules.module.Module, Union[Tuple[oneflow.Tensor, …], oneflow.Tensor], Union[Tuple[oneflow.Tensor, …], oneflow.Tensor]], Union[None, oneflow.Tensor]])¶ Registers a backward hook on the module.
The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature:
hook(module, grad_input, grad_output) -> TensorTuple or None
The
grad_input
andgrad_output
areoneflow.TensorTuple
that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place ofgrad_input
in subsequent computations.grad_input
will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries ingrad_input
andgrad_output
will beNone
for all non-Tensor arguments.For technical reasons, when this hook is applied to a Module, its forward function will receive a view of each Tensor passed to the Module. Similarly the caller will receive a view of each Tensor returned by the Module’s forward function.
Warning
Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.
- Returns
a handle that can be used to remove the added hook by calling
handle.remove()
- Return type
oneflow.utils.hooks.RemovableHandle