Hi, thanks for PR!

Nevertheless, working with checkpoint indices, we also need to logs correctness. In your example, best model saved as train.1.pth, that means checkpoint at 1 epoch.... but from logs perspective it's already 2nd epoch – a little bit confusing :)

You are right. I missed that. Since all the logs are epoch zero based, I propose that we keep it as is but update the total number of epochs in the console progress bar. To keep it consistent.

Let me know if you would prefer a full refactoring of the epoch number so that it starts at 1 everywhere & I will do that instead.

Conceptually, I think it's much better to make full refactoring and save model after N epoch as checkpoint.N.pth. It's much easier to understand :)

Looks like you need to refactor here and here... at least I hope so :)
We still need to check logs, metrics and overall correctness.

@Scitator let me know what you think of this:

Added epoch_log & stage_epoch_log properties to RunnerState to keep it cleaner.

Multi-Stage Training Example

Stage 1:

[...]

Stage 2:

Loading from Stage 2 Checkpoint:

Looking at the plots

Stage 1:

Stage 2:

Loading from Stage 2 Checkpoint:

Was not expecting travis to fail for such a small change. Let me look into it first. Sorry.

Hah, @ostamand looks like you also need to rewrite tests a bit :)
We have check like train.1.loss < train.0.loss.... but with your new indexing... we need +1 there

@Scitator Cool thanks for the tests commit! I was about to take a look at it. Let me know if there is any other change you want me to do?