computation_traces.append(grad_transform_on_trace(computation_trc))

It'd be reasonable to append the output of `grad_transform_on_trace` to `computation_traces: list[TraceCtx]`

Author: crcrpar

thunder/__init__.py

@@ -525,6 +525,7 @@ def apply_transforms_and_build_cache_entry(cd, cs, cache_info, prologue_trc, com
                 from thunder.transforms.autodiff import grad_transform_on_trace
 
                 computation_trc = grad_transform_on_trace(computation_trc)
+                computation_traces.append(computation_trc)
 
             from thunder.executors.passes import _transform_for_operator_executor_execution
             from thunder.distributed.utils import maybe_sort_waits