zero_grad

UNet.zero_grad(set_to_none: bool = True) → None

Resets gradients of all model parameters. See similar function under torch.optim.Optimizer for more context.

Parameters:

set_to_none (bool) – instead of setting to zero, set the grads to None. See torch.optim.Optimizer.zero_grad() for details.