fix epoch number in log (#868)

Summary:
Pull Request resolved: #868

This would currently log "Completed N steps in train epoch 0" in the first epoch, and all epochs would be offset by 1.

Reviewed By: anshulverma, JKSenthil

Differential Revision: D59976202

fbshipit-source-id: 9ef12f3ec86db52d063788e881d9f9d3f1209b7e

Author: galrotem

torchtnt/framework/train.py

@@ -258,8 +258,9 @@ def _train_epoch_impl(
     if not any_steps_completed:
         logger.warning("No steps completed during train epoch!")
     else:
+        # increment epoch happens after we log, hence we add 1 to the epoch counter
         logger.info(
-            f"Completed {train_unit.train_progress.num_steps_completed_in_epoch} steps in train epoch {train_unit.train_progress.num_epochs_completed}"
+            f"Completed {train_unit.train_progress.num_steps_completed_in_epoch} steps in train epoch {train_unit.train_progress.num_epochs_completed + 1}"
         )
 
     # set progress counters for the next epoch