diff --git a/src/model/trainer.py b/src/model/trainer.py
index 02f8449..c855c0f 100644
--- a/src/model/trainer.py
+++ b/src/model/trainer.py
@@ -1042,7 +1042,7 @@ def train(
         model=model,
         train_dataloader=train_dataloader,
         eval_dataloader=eval_dataloader,
-        total_steps=int(max_iter_length / batch_size),
+        total_steps=int(max_iter_length * num_epochs / batch_size),
         output_dir=output_dir,
         num_epochs=num_epochs,
         learning_rate=learning_rate,