Fix minor issues

maestro/trainer/common/utils/leaderboard.py

@@ -25,6 +25,7 @@ def register_checkpoint(self, epoch: int, path: str, loss: float) -> Tuple[bool,
         if loss >= max_loss_in_leaderboard:  # type: ignore
             return False, None
         to_be_removed, _ = self._leaderboard.pop(max_loss_key)  # type: ignore
+        self._leaderboard[epoch] = (path, loss)
         return True, to_be_removed
 
     def get_best_model(self) -> str:

maestro/trainer/models/florence_2/training.py

@@ -208,7 +208,7 @@ def run_training_loop(
             model=model,
             train_loader=train_loader,
             val_loader=val_loader,
-            epoch_number=epoch,
+            epoch_number=epoch + 1,
             configuration=configuration,
             optimizer=optimizer,
             lr_scheduler=lr_scheduler,
@@ -233,7 +233,7 @@ def run_training_epoch(
 ) -> None:
     model.train()
     training_losses: List[float] = []
-    training_iterator = tqdm(train_loader, desc=f"Epoch {epoch_number + 1}/{configuration.training_epochs}")
+    training_iterator = tqdm(train_loader, desc=f"Epoch {epoch_number}/{configuration.training_epochs}")
     for step_id, (inputs, answers) in enumerate(training_iterator):
         input_ids = inputs["input_ids"]
         pixel_values = inputs["pixel_values"]
@@ -257,7 +257,7 @@ def run_training_epoch(
         last_100_losses = training_losses[-100:]
         loss_moving_average = sum(last_100_losses) / len(last_100_losses) if len(last_100_losses) > 0 else 0.0
         training_iterator.set_description(
-            f"Epoch {epoch_number + 1}/{configuration.training_epochs}. Loss: {round(loss_moving_average, 4)}"
+            f"Epoch {epoch_number}/{configuration.training_epochs}. Loss: {round(loss_moving_average, 4)}"
         )
     if len(training_losses) > 0:
         avg_train_loss = sum(training_losses) / len(training_losses)