oneflow.optim.Optimizer.zero_grad

Optimizer.zero_grad(set_to_none: bool = False)

Sets the gradients of all optimized oneflow.Tensor s to zero.

Parameters

set_to_none (bool) – instead of setting to zero, set the grads to None. This will in general have lower memory footprint, and can modestly improve performance. However, it changes certain behaviors.

For example:

1. When the user tries to access a gradient and perform manual ops on it, a None attribute or a Tensor full of 0s will behave differently.

2. If the user requests zero_grad(set_to_none=True) followed by a backward pass, grads are guaranteed to be None for params that did not receive a gradient.

3. Optimizers have a different behavior if the gradient is 0 or None (in one case it does the step with a gradient of 0 and in the other it skips the step altogether).