changed the file place and deleted unnecessary code

cehongwang · cehongwang · commit a0552c5310e5 · 2025-03-26T07:58:18.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
@@ -734,10 +734,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
`@@ -734,10 +734,9 @@ def run(`
`734`	`734`	`builder_config, self.compilation_settings.timing_cache_path`
`735`	`735`	`)`
`736`	`736`	`# TODO: Memory control prototyping. Under discussion`
`737`		`- self.module.to("cpu")`
`738`		`- torch.cuda.empty_cache()`
`739`	`737`	`del self.module`
`740`	`738`	`gc.collect()`
	`739`	`+ torch.cuda.empty_cache()`
`741`	`740`	`serialized_engine = self.builder.build_serialized_network(`
`742`	`741`	`self.ctx.net, builder_config`
`743`	`742`	`)`