Fixed Cuda Error

Author: cehongwang

py/torch_tensorrt/dynamo/_refit.py

@@ -109,7 +109,7 @@ def construct_refit_mapping(
 
 
 def construct_refit_mapping_from_weight_name_map(
-    weight_name_map: dict[Any, Any], state_dict: dict[Any, Any]
+    weight_name_map: dict[Any, Any], state_dict: dict[Any, Any], device: torch.device
 ) -> dict[Any, Any]:
     engine_weight_map = {}
     for engine_weight_name, (sd_weight_name, np_weight_type) in weight_name_map.items():
@@ -120,7 +120,11 @@ def construct_refit_mapping_from_weight_name_map(
             # If weights is not in sd, we can leave it unchanged
             continue
         else:
-            engine_weight_map[engine_weight_name] = state_dict[sd_weight_name]
+            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")
+            )
 
         engine_weight_map[engine_weight_name] = (
             engine_weight_map[engine_weight_name]
@@ -162,7 +166,7 @@ def _refit_single_trt_engine_with_gm(
             "constant_mapping", {}
         )  # type: ignore
         mapping = construct_refit_mapping_from_weight_name_map(
-            weight_name_map, new_gm.state_dict()
+            weight_name_map, new_gm.state_dict(), torch_device
         )
         constant_mapping_with_type = {}