tabensemb.model.AbstractNN.optimizer_zero_grad

method

AbstractNN.optimizer_zero_grad(epoch: int, batch_idx: int, optimizer: Optimizer, optimizer_idx: int) → None

Override this method to change the default behaviour of optimizer.zero_grad().

Parameters:

• epoch – Current epoch

• batch_idx – Index of current batch

• optimizer – A PyTorch optimizer

• optimizer_idx – If you used multiple optimizers this indexes into that list.

Examples:

# DEFAULT
def optimizer_zero_grad(self, epoch, batch_idx, optimizer, optimizer_idx):
    optimizer.zero_grad()

# Set gradients to `None` instead of zero to improve performance (not required on `torch>=2.0.0`).
def optimizer_zero_grad(self, epoch, batch_idx, optimizer, optimizer_idx):
    optimizer.zero_grad(set_to_none=True)

See torch.optim.Optimizer.zero_grad() for the explanation of the above example.