detach_state_graph

torch_sim.state.detach_state_graph(state)

Detach any autograd-graph-carrying tensors on a state, in place.

Some models return graph-carrying outputs - notably UMA, whose energy keeps requires_grad=True even though its forces are already detached. When a converged system is popped out of the running batch and accumulated for the remainder of the run, such a tensor pins that swap’s entire forward autograd graph, so live GPU memory grows by roughly one graph per finished system until the device fills - independent of the batch size and unreclaimable by empty_cache (it is allocated, not cached). Completed states are only ever read for their values, never differentiated, so dropping the graph is safe.

Parameters:

state (SimState) – State whose grad-carrying tensor attributes are detached in place.

Returns:

The same state instance, with grad-carrying tensors detached.

Return type:

SimState