Skip to content

Commit 6d36077

Browse files
committed
changed the file place and deleted unnecessary code
1 parent c698e60 commit 6d36077

File tree

2 files changed

+2
-2
lines changed

2 files changed

+2
-2
lines changed

demo/flux_demo.py examples/apps/flux_demo.py

+1
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@
8888
del ep
8989
pipe.transformer = trt_gm
9090
pipe.transformer.config = config
91+
trt_gm.device = torch.device("cuda")
9192
torch.cuda.empty_cache()
9293

9394

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -711,10 +711,9 @@ def run(
711711
builder_config, self.compilation_settings.timing_cache_path
712712
)
713713
# TODO: Memory control prototyping. Under discussion
714-
self.module.to("cpu")
715-
torch.cuda.empty_cache()
716714
del self.module
717715
gc.collect()
716+
torch.cuda.empty_cache()
718717
serialized_engine = self.builder.build_serialized_network(
719718
self.ctx.net, builder_config
720719
)

0 commit comments

Comments
 (0)