fix refit issue for torch.compile

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -40,7 +40,7 @@
     get_node_name,
     get_trt_tensor,
 )
-from torch_tensorrt.dynamo.utils import DYNAMIC_DIM, to_torch_device
+from torch_tensorrt.dynamo.utils import DYNAMIC_DIM, get_model_device, to_torch_device
 from torch_tensorrt.fx.observer import Observer
 from torch_tensorrt.logging import TRT_LOGGER
 
@@ -434,9 +434,8 @@ def _save_weight_mapping(self) -> None:
         """
         _LOGGER.info("Building weight name mapping...")
         # Stage 1: Name mapping
-        sd = self.module.state_dict()
         torch_device = to_torch_device(self.compilation_settings.device)
-        gm_is_on_cuda = list(sd.values())[0].device.type == "cuda"
+        gm_is_on_cuda = get_model_device(self.module).type == "cuda"
         if not gm_is_on_cuda:
             # If the model original position is on CPU, move it GPU
             sd = {

tests/py/dynamo/models/test_engine_cache.py

@@ -184,9 +184,6 @@ def test_dynamo_compile_with_custom_engine_cache(self):
             msg=f"Engine caching didn't speed up the compilation. Time taken without engine caching: {times[0]} ms, time taken with engine caching: {times[2]} ms",
         )
 
-    @pytest.mark.skip(
-        reason="The test needs a fix for refit, which is reported in https://github.com/pytorch/TensorRT/issues/3126"
-    )
     def test_torch_compile_with_default_disk_engine_cache(self):
         # Custom Engine Cache
         model = models.resnet18(pretrained=True).eval().to("cuda")
@@ -251,9 +248,6 @@ def test_torch_compile_with_default_disk_engine_cache(self):
             msg=f"Engine caching didn't speed up the compilation. Time taken without engine caching: {times[0]} ms, time taken with engine caching: {times[2]} ms",
         )
 
-    @pytest.mark.skip(
-        reason="The test needs a fix for refit, which is reported in https://github.com/pytorch/TensorRT/issues/3126"
-    )
     def test_torch_compile_with_custom_engine_cache(self):
         # Custom Engine Cache
         model = models.resnet18(pretrained=True).eval().to("cuda")