fix: detach tensors

metaopt · XuehaiPan · Jan 15, 2023 · Nov 16, 2022 · Nov 18, 2022 · Nov 18, 2022
commit 7b3d708b45e2f7630f888a3714567c3ffbd93bcf
diff --git a/torchopt/diff/implicit/decorator.py b/torchopt/diff/implicit/decorator.py
@@ -289,6 +289,7 @@ def forward(  # type: ignore[override] # pylint: disable=arguments-differ
                         f'solver_fn should be a torch.Tensor or a tuple of torch.Tensor. '
                         f'Got {output}'
                     )
+                output_detached = tuple(t.data for t in output)
 
                 (
                     args_treespec,
@@ -300,7 +301,7 @@ def forward(  # type: ignore[override] # pylint: disable=arguments-differ
                 ctx.args_is_tensor_mask = args_is_tensor_mask
                 ctx.args_non_tensors = args_non_tensors
 
-                ctx.save_for_backward(*output, *args_tensors)
+                ctx.save_for_backward(*output_detached, *args_tensors)
                 ctx.output_is_tensor = output_is_tensor
 
                 return (*output, aux, output_is_tensor, type(output))

diff --git a/torchopt/diff/implicit/nn/module.py b/torchopt/diff/implicit/nn/module.py
@@ -169,7 +169,17 @@ def solver_fn(
             **kwargs,
         ) -> Tuple[TupleOfTensors, Any]:
             output = cls_solve(self, *input, **kwargs)
-            flat_optimal_params: TupleOfTensors = tuple(pytree.tree_leaves(params_containers))  # type: ignore[arg-type]
+
+            def detach_(t):
+                requires_grad = t.requires_grad
+                return t.detach_().requires_grad_(requires_grad)
+
+            flat_optimal_params: TupleOfTensors = tuple(
+                map(
+                    detach_,
+                    pytree.tree_leaves(params_containers),  # type: ignore[arg-type]
+                )
+            )
             return flat_optimal_params, output
 
         # pylint: disable-next=unused-variable