changed the file place and deleted unnecessary code

cehongwang · cehongwang · commit b2eb297eb73d · 2025-02-28T23:18:11.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
@@ -706,10 +706,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
`@@ -706,10 +706,9 @@ def run(`
`706`	`706`	`builder_config, self.compilation_settings.timing_cache_path`
`707`	`707`	`)`
`708`	`708`	`# TODO: Memory control prototyping. Under discussion`
`709`		`- self.module.to("cpu")`
`710`		`- torch.cuda.empty_cache()`
`711`	`709`	`del self.module`
`712`	`710`	`gc.collect()`
	`711`	`+ torch.cuda.empty_cache()`
`713`	`712`	`serialized_engine = self.builder.build_serialized_network(`
`714`	`713`	`self.ctx.net, builder_config`
`715`	`714`	`)`