Always enable grad before connecting to autograd

Author: shino16

thunder/__init__.py

@@ -865,13 +865,14 @@ def fn_(*args, **kwargs) -> Any:
         cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
 
         result = cache_entry.computation_fn(*inps)
-        # We must do this before connecting to autograd so that
-        # grad_fn attribute will be set on the outputs accordingly
-        pytorch.set_grad_enabled(cd.is_grad_enabled)
         result = maybe_connect_to_autograd(cache_entry, result)
         result = call_epilogue(cache_entry, result, pro_to_epi)
 
         cs.last_computation = cache_entry.computation_fn
+
+        # Reflect the state of is_grad_enabled, as its changes have been recorded only inside Thunder
+        pytorch.set_grad_enabled(cd.is_grad_enabled)
+
         return result
 
     if isinstance(fn, pytorch.nn.Module):

thunder/core/symbol.py

@@ -329,6 +329,7 @@ def __call__(self, *args, **kwargs):
         if cd is not None and not cd.is_grad_enabled:
             flat_args, _ = tree_flatten((args, kwargs))
             flat_arg_ids = {id(arg) for arg in flat_args}
+
             # If grad is disabled using `torch.no_grad` or `torch._C._set_grad_enabled(False)`,
             # tag the results with `DETACHED_AUTOGRAD_GRAPH` which makes this Symbol a constant for
             # vjp transform (applied later).

thunder/executors/torch_autograd.py

@@ -168,16 +168,17 @@ def connect_to_autograd(
             disable_split_autograd, lambda: "is_differentiable_outputs is not supported when split_autograd is enabled"
         )
 
-    dummy_res = ThunderFunction.apply(
-        return_none_instead_of_grads,
-        backward_fn,
-        side_channel,
-        saved_tensors,
-        saved_other,
-        is_differentiable_outputs,
-        flat_output,
-        *flat_args,
-    )
+    with torch.enable_grad():
+        dummy_res = ThunderFunction.apply(
+            return_none_instead_of_grads,
+            backward_fn,
+            side_channel,
+            saved_tensors,
+            saved_other,
+            is_differentiable_outputs,
+            flat_output,
+            *flat_args,
+        )
     if side_channel is not None:
         # we need to pass the inputs to avoid "leave has moved inside the graph"
         # if the function returns an argument as is