oneflow.Tensor.register_hook

Tensor.register_hook(hook)

Registers a backward hook.

The hook will be called every time a gradient with respect to the Tensor is computed. The hook should have the following signature:

hook(grad) -> Tensor or None

The hook should not modify its argument, but it can optionally return a new gradient which will be used in place of grad.

For example:

>>> import oneflow as flow
>>> x = flow.ones(5, requires_grad=True)
>>> def hook(grad):
...     return grad * 2
>>> x.register_hook(hook)
>>> y = x * 2
>>> y.sum().backward()
>>> x.grad
tensor([4., 4., 4., 4., 4.], dtype=oneflow.float32)