Machine calls cleanup on time-step cleanup (#536)

ihmeuw · Nov 13, 2024 · 2a81a36 · 2a81a36
1 parent ed09d19
commit 2a81a36
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/vivarium/framework/state_machine.py b/src/vivarium/framework/state_machine.py
@@ -563,6 +563,9 @@ def on_initialize_simulants(self, pop_data: SimulantData) -> None:
     def on_time_step(self, event: Event) -> None:
         self.transition(event.index, event.time)
 
+    def on_time_step_cleanup(self, event: Event) -> None:
+        self.cleanup(event.index, event.time)
+
     ##################
     # Public methods #
     ##################