add new states when the max_memory_scaler is updated (#222)

kianpu34593 · web-flow · commit 5371cb6a83cd · 2025-07-24T23:17:14.000-04:00
diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py
@@ -952,6 +952,8 @@ def _get_first_batch(self) -> SimState:
                 scale_factor=self.memory_scaling_factor,
             )
             self.max_memory_scaler = self.max_memory_scaler * self.max_memory_padding
+            newer_states = self._get_next_states()
+            states = [*states, *newer_states]
         return concatenate_states([first_state, *states])
 
     def next_batch(  # noqa: C901

Original file line number	Diff line number	Diff line change
`@@ -952,6 +952,8 @@ def _get_first_batch(self) -> SimState:`
`952`	`952`	`scale_factor=self.memory_scaling_factor,`
`953`	`953`	`)`
`954`	`954`	`self.max_memory_scaler = self.max_memory_scaler * self.max_memory_padding`
	`955`	`+ newer_states = self._get_next_states()`
	`956`	`+ states = [states, newer_states]`
`955`	`957`	`return concatenate_states([first_state, *states])`
`956`	`958`
`957`	`959`	`def next_batch( # noqa: C901`