Increment epochs based on last_batch() instead of at the end of the train loop.

[justinxzhao]

self._train_loop() trains over a full epoch of data or up to the last training step, whichever is sooner, so we should only increment epochs when last_batch() is True as a check within self._train_loop().

This fixes a potential issue with the last epoch being 1 too high when the training runway concludes partway through an epoch.

[github-actions]

Unit Test Results

       5 files   -   1         5 suites   - 1   51m 6s ⏱️ -5s
2 807 tests ±  0  2 793 ✔️ ±0  12 💤 ±0  2 ❌ ±0 
2 835 runs   - 12  2 815 ✔️  - 9  18 💤  - 3  2 ❌ ±0 

For more details on these failures, see this check.

Results for commit e7178ab. ± Comparison against base commit 4af5331.

♻️ This comment has been updated with latest results.