🐛 Fix minor bugs

hezarai · Jun 3, 2024 · a5652ef · a5652ef
1 parent cfec53c
commit a5652ef
Showing 1 changed file with 4 additions and 1 deletion.
diff --git 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)