🐛 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)