Finalize the refit revision

Author: cehongwang

examples/apps/flux-demo.py

@@ -41,7 +41,7 @@
     "use_fp32_acc": True,
     "use_explicit_typing": True,
     "debug": False,
-    "use_python_runtime": False,
+    "use_python_runtime": True,
     "immutable_weights": False,
     # "cache_built_engines": True,
     # "reuse_cached_engines": True,

examples/dynamo/torch_export_flux_dev.py

@@ -112,6 +112,7 @@
     min_block_size=1,
     use_fp32_acc=True,
     use_explicit_typing=True,
+    use_python_runtime=True,
 )
 
 # %%
@@ -126,7 +127,7 @@
 torch.cuda.empty_cache()
 pipe.transformer = trt_gm
 pipe.transformer.config = config
-
+trt_gm.device = torch.device("cuda")
 # %%
 # Image generation using prompt
 # ---------------------------

py/torch_tensorrt/dynamo/_refit.py

@@ -507,23 +507,22 @@ def refit_module_weights(
         serialization_config.clear_flag(trt.SerializationFlag.EXCLUDE_WEIGHTS)
         serialized_engine = engine.serialize_with_config(serialization_config)
 
-        del engine
-        gc.collect()
-        torch.cuda.empty_cache()
-
-        if isinstance(
-            compiled_submodule, (PythonTorchTensorRTModule, TorchTensorRTModule)
-        ):
+        if isinstance(compiled_submodule, PythonTorchTensorRTModule):
+            compiled_submodule.serialized_engine = bytes(serialized_engine)
+        elif isinstance(compiled_submodule, TorchTensorRTModule):
             compiled_submodule.engine = None  # Clear the engine for TorchTensorRTModule, otherwise it won't be updated
             compiled_submodule.serialized_engine = bytes(serialized_engine)
             compiled_submodule.setup_engine()
-
         elif inline_module:
             new_engine_info = list(engine_info)
             new_engine_info[ENGINE_IDX] = bytes(serialized_engine)
             refitted_engine = torch.classes.tensorrt.Engine(tuple(new_engine_info))
             setattr(compiled_module, f"{name}_engine", refitted_engine)
 
+        del engine
+        gc.collect()
+        torch.cuda.empty_cache()
+
     # TODO: Memory control prototyping. Under discussion
     if settings.offload_module_to_cpu:
         del new_partitioned_module

tests/py/dynamo/models/test_model_refit.py

@@ -763,6 +763,7 @@ def forward(self, x):
             debug=True,
             min_block_size=1,
             immutable_weights=False,
+            offload_module_to_cpu=False,
         )
 
     num_pyt_segments = len(