From a5652efb28c88dff2ab5162e6fb7042e5a065b6f Mon Sep 17 00:00:00 2001 From: arxyzan Date: Mon, 3 Jun 2024 14:48:02 +0330 Subject: [PATCH] :bug: Fix minor bugs --- hezar/trainer/trainer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/hezar/trainer/trainer.py b/hezar/trainer/trainer.py index abe9f050..c47cbd48 100644 --- a/hezar/trainer/trainer.py +++ b/hezar/trainer/trainer.py @@ -538,7 +538,10 @@ def inner_training_loop(self, epoch_num: int): # Gather outputs for metrics losses_sum += outputs["loss"].item() - if self.state.epoch_step % self.config.gradient_accumulation_steps == 0: + if ( + self.state.epoch_step % self.config.gradient_accumulation_steps == 0 + or self.state.epoch_step == self.steps_in_epoch + ): avg_loss = losses_sum * self.config.gradient_accumulation_steps / self.state.epoch_step iterator.set_postfix(loss=avg_loss)