How does optax.MultiSteps handle the last few steps?

[ayaka14732]

I am using optax.MultiSteps for gradient accumulation.

My question is: if my n_accumulation_steps is 8, total_step is 807, then 807 % 8 == 7, there would be 7 steps left not updated in the end the epoch.

Is there anything like optimizer.finalize() that can update the last few steps?