diff --git a/side_effects_penalties/side_effects_penalty_test.py b/side_effects_penalties/side_effects_penalty_test.py
index 2468216..69dc104 100644
--- a/side_effects_penalties/side_effects_penalty_test.py
+++ b/side_effects_penalties/side_effects_penalty_test.py
@@ -67,6 +67,8 @@ class BaselineTestCase(SideEffectsTestCase):
       }
       self.assertEqual(baseline_state, comparison_dict[key])
       current_state = next_state
+      if timestep.last():
+        return
 
 
 class StartBaselineTest(BaselineTestCase):