diff --git a/src/post_training/callbacks/inference_checkpoint.py b/src/post_training/callbacks/inference_checkpoint.py
index 658b677..f95df75 100644
--- a/src/post_training/callbacks/inference_checkpoint.py
+++ b/src/post_training/callbacks/inference_checkpoint.py
@@ -54,7 +54,10 @@ def on_step_end(
 
         if state.global_step == 0:
             return
-        if state.global_step % self.save_steps != 0:
+
+        is_periodic = state.global_step % self.save_steps == 0
+        is_last_step = state.global_step == state.max_steps
+        if not (is_periodic or is_last_step):
             return
 
         # Only save on the main process.