Skip to content

Commit bd9c29a

Browse files
authored
fix: Repair null bindings issue in TRT Engines (#2080)
1 parent 6ceaed8 commit bd9c29a

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

py/torch_tensorrt/dynamo/backend/backends.py

+7-1
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,9 @@ def _compile_module(
121121
torch_executed_ops=settings.torch_executed_ops,
122122
)
123123

124+
# Store TRT replicas of Torch subgraphs
125+
trt_modules = {}
126+
124127
# Iterate over all components that can be accelerated
125128
# Generate the corresponding TRT Module for those
126129
for name, _ in partitioned_module.named_children():
@@ -138,7 +141,10 @@ def _compile_module(
138141
settings=settings,
139142
)
140143

141-
# Replace FX Module with TRT Module
144+
trt_modules[name] = trt_mod
145+
146+
# Replace all FX Modules with TRT Modules
147+
for name, trt_mod in trt_modules.items():
142148
setattr(partitioned_module, name, trt_mod)
143149

144150
return partitioned_module

0 commit comments

Comments
 (0)