fix(optim): move apply_updates in grad context

Author: XuehaiPan

torchopt/optim/func/base.py

@@ -90,14 +90,13 @@ def step(
         with torch.enable_grad():
             # Step parameters only
             grads = torch.autograd.grad(loss, params, create_graph=True, allow_unused=True)
-
-        updates, self.optim_state = self.impl.update(
-            grads,
-            self.optim_state,
-            params=params,
-            inplace=inplace,
-        )
-        return apply_updates(params, updates, inplace=inplace)
+            updates, self.optim_state = self.impl.update(
+                grads,
+                self.optim_state,
+                params=params,
+                inplace=inplace,
+            )
+            return apply_updates(params, updates, inplace=inplace)
 
     def state_dict(self) -> OptState:
         """Extract the references of the optimizer states.

torchopt/optim/meta/base.py

@@ -78,10 +78,14 @@ def step(self, loss: torch.Tensor) -> None:  # pylint: disable=too-many-locals
             with torch.enable_grad():
                 # Step parameters only
                 grads = torch.autograd.grad(loss, flat_params, create_graph=True, allow_unused=True)
+                updates, new_state = self.impl.update(
+                    grads,
+                    state,
+                    params=flat_params,
+                    inplace=False,
+                )
+                flat_new_params = apply_updates(flat_params, updates, inplace=False)
 
-            updates, new_state = self.impl.update(grads, state, params=flat_params, inplace=False)
-
-            flat_new_params = apply_updates(flat_params, updates, inplace=False)
             new_params: ModuleTensorContainers = pytree.tree_unflatten(  # type: ignore[assignment]
                 container_treespec,
                 flat_new_params,