kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-14 06:27:45 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None] [AutoDeploy] canonicalize_graph before shape prop for consistent state_dict (#7223)

Browse Source 
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>
This commit is contained in:
Lucas Liebenwein 2025-08-25 13:59:57 -07:00 committed by GitHub
parent bf1b958f1a
commit 97d550b4ba
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorrt_llm/_torch/auto_deploy/transform/interface.py
View File

@ -326,6 +326,7 @@ class BaseTransform(ABC):
# check if run cleanup depending on the config and info
if self.config.requires_shape_prop and not has_valid_shapes:
canonicalize_graph(gm)
with lift_to_meta(gm):
canonicalize_graph(gm, shape_prop=True)
is_clean = True
@ -351,6 +352,7 @@ class BaseTransform(ABC):
# check if run cleanup depending on the config and info
if self.config.run_shape_prop and not (info.is_clean and info.has_valid_shapes):
canonicalize_graph(gm)
with lift_to_meta(gm):
canonicalize_graph(gm, shape_prop=True)
elif self.config.run_graph_cleanup and not info.is_clean: