changed the file place and deleted unnecessary code

cehongwang · cehongwang · commit 6d36077d31da · 2025-02-28T23:19:54.000Z
diff --git a/examples/apps/flux_demo.py b/examples/apps/flux_demo.py
@@ -88,6 +88,7 @@
 del ep
 pipe.transformer = trt_gm
 pipe.transformer.config = config
+trt_gm.device = torch.device("cuda")
 torch.cuda.empty_cache()
 
 
diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py
@@ -711,10 +711,9 @@ def run(
             builder_config, self.compilation_settings.timing_cache_path
         )
         # TODO: Memory control prototyping. Under discussion
-        self.module.to("cpu")
-        torch.cuda.empty_cache()
         del self.module
         gc.collect()
+        torch.cuda.empty_cache()
         serialized_engine = self.builder.build_serialized_network(
             self.ctx.net, builder_config
         )

Original file line number	Diff line number	Diff line change
`@@ -711,10 +711,9 @@ def run(`
`711`	`711`	`builder_config, self.compilation_settings.timing_cache_path`
`712`	`712`	`)`
`713`	`713`	`# TODO: Memory control prototyping. Under discussion`
`714`		`- self.module.to("cpu")`
`715`		`- torch.cuda.empty_cache()`
`716`	`714`	`del self.module`
`717`	`715`	`gc.collect()`
	`716`	`+ torch.cuda.empty_cache()`
`718`	`717`	`serialized_engine = self.builder.build_serialized_network(`
`719`	`718`	`self.ctx.net, builder_config`
`720`	`719`	`)`