Optimizer step hook missing param
Version
5.2.1
Describe the bug
TypeError: optimizer_step() takes 5 positional arguments but 6 were given
To Reproduce
Overload optimizer_step as outlined in docs
def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i):
The Source
line in trainer gives an extra optimizer_closure param, triggering the bug
model.optimizer_step(self.current_epoch, batch_nb, optimizer, opt_idx, optimizer_closure)
To Fix
adding a 6th argument to the definition fixes the issue, not sure if this is intended and only a change in docs is required
def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_i, closure):