Do not tag argument proxy as DETACHED_AUTOGRAD_GRAPH

Author: shino16

thunder/core/symbol.py

@@ -327,11 +327,13 @@ def __call__(self, *args, **kwargs):
 
         cd = get_compile_data()
         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).
             def tag_tensorproxy_output_as_detached(proxy):
-                if isinstance(proxy, TensorProxy):
+                if isinstance(proxy, TensorProxy) and id(proxy) not in flat_arg_ids:
                     proxy.tags.add(ProxyTag.DETACHED_AUTOGRAD_GRAPH)
 
                 return proxy