File tree Expand file tree Collapse file tree 1 file changed +7
-1
lines changed
py/torch_tensorrt/dynamo/backend Expand file tree Collapse file tree 1 file changed +7
-1
lines changed Original file line number Diff line number Diff 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
You can’t perform that action at this time.
0 commit comments