fix optimizer runner too

orionarcher · orionarcher · commit cdf605044ca9 · 2025-03-04T22:27:27.000Z
diff --git a/torchsim/runners.py b/torchsim/runners.py
@@ -136,11 +136,11 @@ def convergence_fn(state: BaseState, last_energy: torch.Tensor) -> bool:
         )
     state: BaseState = initialize_state(system, model.device, model.dtype)
 
-    state, update_fn = optimizer(
-        state=state,
+    init_fn, update_fn = optimizer(
         model=model,
         **optimizer_kwargs,
     )
+    state = init_fn(state)
 
     step: int = 1
     last_energy = state.energy + 1

Original file line number	Diff line number	Diff line change
`@@ -136,11 +136,11 @@ def convergence_fn(state: BaseState, last_energy: torch.Tensor) -> bool:`
`136`	`136`	`)`
`137`	`137`	`state: BaseState = initialize_state(system, model.device, model.dtype)`
`138`	`138`
`139`		`- state, update_fn = optimizer(`
`140`		`- state=state,`
	`139`	`+ init_fn, update_fn = optimizer(`
`141`	`140`	`model=model,`
`142`	`141`	`**optimizer_kwargs,`
`143`	`142`	`)`
	`143`	`+ state = init_fn(state)`
`144`	`144`
`145`	`145`	`step: int = 1`
`146`	`146`	`last_energy = state.energy + 1`