Fixed the bug of SDXL Cuda Error

Author: cehongwang

py/torch_tensorrt/dynamo/_refit.py

@@ -123,7 +123,7 @@ def construct_refit_mapping_from_weight_name_map(
             engine_weight_map[engine_weight_name] = (
                 state_dict[sd_weight_name]
                 if state_dict[sd_weight_name].device == device
-                else state_dict[sd_weight_name].to("device")
+                else state_dict[sd_weight_name].to(device)
             )
 
         engine_weight_map[engine_weight_name] = (

py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py

@@ -7,6 +7,7 @@
 import numpy as np
 import torch
 import torch_tensorrt
+from torch.export._trace import _export
 from torch_tensorrt._Device import Device
 from torch_tensorrt.dynamo import _defaults
 from torch_tensorrt.dynamo._compiler import compile as dynamo_compile
@@ -309,7 +310,7 @@ def refit_gm(self) -> None:
 
     def get_exported_program(self) -> torch.export.ExportedProgram:
         if self.allow_complex_guards_as_runtime_asserts:
-            return torch.export._trace._export(
+            return _export(
                 self.original_model,
                 self.arg_inputs,
                 kwargs=self.kwarg_inputs,