From 8a1d55a9db59be79416608ad116ec5a36cd4ceb4 Mon Sep 17 00:00:00 2001
From: Luis Pineda <4759586+luisenp@users.noreply.github.com>
Date: Fri, 16 Jun 2023 09:20:45 -0700
Subject: [PATCH] Add autograd create graph and strategy options for
 AutoDiffCostFunction.

---
 theseus/core/cost_function.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/theseus/core/cost_function.py b/theseus/core/cost_function.py
index 186d1e65b..29482aa7d 100644
--- a/theseus/core/cost_function.py
+++ b/theseus/core/cost_function.py
@@ -209,8 +209,10 @@ def __init__(
         cost_weight: Optional[CostWeight] = None,
         aux_vars: Optional[Sequence[Variable]] = None,
         name: Optional[str] = None,
+        autograd_create_graph: bool = True,
         autograd_strict: bool = False,
         autograd_vectorize: bool = False,
+        autograd_strategy: str = "reverse-mode",
         autograd_mode: Union[str, AutogradMode] = AutogradMode.VMAP,
     ):
         if cost_weight is None:
@@ -228,8 +230,10 @@ def __init__(
 
         self._err_fn = err_fn
         self._dim = dim
+        self._autograd_create_graph = autograd_create_graph
         self._autograd_strict = autograd_strict
         self._autograd_vectorize = autograd_vectorize
+        self._autograd_strategy = autograd_strategy
 
         # The following are auxiliary Variable objects to hold tensor data
         # during jacobian computation without modifying the original Variable objects
@@ -285,9 +289,10 @@ def _compute_autograd_jacobian(
         return autogradF.jacobian(
             jac_fn,
             optim_tensors,
-            create_graph=True,
+            create_graph=self._autograd_create_graph,
             strict=self._autograd_strict,
             vectorize=self._autograd_vectorize,
+            strategy=self._autograd_strategy,
         )
 
     def _make_jac_fn_vmap(