diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py
index 902e5a8d..bbad57cb 100644
--- a/src/somd2/runner/_runner.py
+++ b/src/somd2/runner/_runner.py
@@ -656,9 +656,16 @@ def generate_lam_vals(lambda_base, increment=0.001):
 
             # Get the simulation speed.
             speed = dynamics.time_speed()
-
             # Checkpoint.
-            self._checkpoint(system, index, 0, speed, is_final_block=True)
+            self._checkpoint(
+                system,
+                index,
+                0,
+                speed,
+                is_final_block=True,
+                lambda_grad=lambda_grad,
+                lambda_energy=lambda_energy,
+            )
 
             _logger.success(
                 f"{_lam_sym} = {lambda_value:.5f} complete, speed = {speed:.2f} ns day-1"