From 6b736369c3bd3f70905e2f350295ecb75e21ce8b Mon Sep 17 00:00:00 2001 From: Masaki Kozuki Date: Fri, 13 Dec 2024 08:21:27 +0900 Subject: [PATCH] add tensor subclass transform output to traces Signed-off-by: Masaki Kozuki --- thunder/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/thunder/__init__.py b/thunder/__init__.py index 3ac277ed4..b16930450 100644 --- a/thunder/__init__.py +++ b/thunder/__init__.py @@ -635,6 +635,7 @@ def get_computation_and_inputs(*args, **kwargs): _tensor_subclass_transform_applied = True if not _tensor_subclass_transform_applied: computation_trc = flatten_tensor_subclasses(computation_trc) + computation_traces.append(computation_trc) if backward_trc is None: from thunder.executors.passes import transform_for_execution as transform_for_execution_pass