tabensemb.model.AbstractNN.on_before_zero_grad

method

AbstractNN.on_before_zero_grad(optimizer: Optimizer) → None

Called after training_step() and before optimizer.zero_grad().

Called in the training loop after taking an optimizer step and before zeroing grads. Good place to inspect weight information with weights updated.

This is where it is called:

for optimizer in optimizers:
    out = training_step(...)

    model.on_before_zero_grad(optimizer) # < ---- called here
    optimizer.zero_grad()

    backward()

Parameters:

optimizer – The optimizer for which grads should be zeroed.