In 1.0.4, training_step is never called when optimizer_step is overriden

[catalys1]

🐛 Bug

Something changed from 1.0.3 to 1.0.4: in 1.0.3 I was able to override optimizer_step in my LightningModule and everything worked fine. But in 1.0.4, when I override optimizer_step it causes training_step to be skipped.

Please reproduce using the BoringModel and post here

BoringModel replication

Expected behavior

Defining optimizer_step should not cause training_step to be skipped over. This is definitely a problem.

[jwspaeth]

Also ran into this bug. Fixed in my case by passing optimizer.step the optimizer_closure, as they do in LightningModule here.

[catalys1]

Thanks, that seems to be working. If this is going to be the required way of handling it (unlike in previous releases) then I think it needs to be clearly pointed out in the documentation. Otherwise, hopefully it can be fixed before the next maintenance update.

[ananyahjha93]

yes, I think the docs are a bit old, I am updating them in the fix. Apart from that just making sure that passing optimizer_closure is the only issue here.

[edenlightning]

@SeanNaren please update what is left TODO here

[SeanNaren]

This should be closed, the docs have been cleared up in #4455, ensuring that users know to pass the optimizer_closure to the optimizer.