diff --git a/training_pipeline/predict.py b/training_pipeline/predict.py
index 27886f7..7b3450b 100644
--- a/training_pipeline/predict.py
+++ b/training_pipeline/predict.py
@@ -92,7 +92,7 @@ def predict_batch(model, image_paths, device="cuda", threshold=0.5, checkpoint_p
             processed_images.append(image_path)
             
             # Save checkpoint every 10 images
-            if checkpoint_path and (i + 1) % 10 == 0:
+            if checkpoint_path and newly_processed % 10 == 0:
                 metadata = {
                     'device': device,
                     'threshold': threshold,