diff --git a/bfps/_code.py b/bfps/_code.py
index 3ff3de341b20c578262c1111835ba799359db1f6..44570240fa1546ea9f9188e8e8f6edf0bd917f07 100644
--- a/bfps/_code.py
+++ b/bfps/_code.py
@@ -32,6 +32,7 @@ import argparse
 import h5py
 from datetime import datetime
 import math
+import warnings
 
 import bfps
 from ._base import _base
diff --git a/bfps/_fluid_base.py b/bfps/_fluid_base.py
index 8ed88ca6ac035717451e5eab9b6fb7709d3bd334..2387d31e972765e0928f671359b437eef7317f74 100644
--- a/bfps/_fluid_base.py
+++ b/bfps/_fluid_base.py
@@ -264,7 +264,7 @@ class _fluid_particle_base(_code):
                                       '<< time_difference/nprocs << " seconds" << std::endl;\n' +
                                   'time0 = time1;\n')
         if not postprocess_mode:
-            self.main       += 'for (int max_iter = iteration+niter_todo; iteration < max_iter; iteration++)\n'
+            self.main       += 'for (int max_iter = iteration+niter_todo-iteration%niter_todo; iteration < max_iter; iteration++)\n'
             self.main       += '{\n'
 
             self.main       += """