fix cross compilation test bug (#3609)

Author: lanluo-nvidia

py/torch_tensorrt/_utils.py

@@ -16,10 +16,11 @@ def check_cross_compile_trt_win_lib() -> bool:
     # cross compile feature is only available on linux
     # build engine on linux and run on windows
     if sys.platform.startswith("linux"):
+        import re
+
         import dllist
 
         loaded_libs = dllist.dllist()
-        target_lib = "libnvinfer_builder_resource_win.so.*"
-        if target_lib in loaded_libs:
-            return True
+        target_lib = ".*libnvinfer_builder_resource_win.so.*"
+        return any(re.match(target_lib, lib) for lib in loaded_libs)
     return False

tests/py/dynamo/runtime/test_003_cross_compile_for_windows.py

@@ -7,8 +7,8 @@
 import torch
 import torch_tensorrt
 from torch.testing._internal.common_utils import TestCase
-from torch_tensorrt.dynamo.utils import get_model_device
 from torch_tensorrt._utils import check_cross_compile_trt_win_lib
+from torch_tensorrt.dynamo.utils import get_model_device
 
 from ..testing_utilities import DECIMALS_OF_AGREEMENT
 
@@ -85,8 +85,12 @@ def forward(self, a, b):
     @pytest.mark.unit
     def test_dynamo_cross_compile_for_windows_cpu_offload(self):
         class Add(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 3)
+
             def forward(self, a, b):
-                return torch.add(a, b)
+                return torch.add(self.linear(a), b)
 
         model = Add().eval().cuda()
         inputs = (torch.randn(2, 3).cuda(), torch.randn(2, 3).cuda())
@@ -101,7 +105,7 @@ def forward(self, a, b):
             trt_gm = torch_tensorrt.dynamo.cross_compile_for_windows(
                 exp_program, **compile_spec
             )
-            assert get_model_device(trt_gm).type == "cpu"
+            assert get_model_device(model).type == "cpu"
             torch_tensorrt.dynamo.save_cross_compiled_exported_program(
                 trt_gm, file_path=trt_ep_path
             )
@@ -112,6 +116,10 @@ def forward(self, a, b):
         platform.system() != "Linux" or platform.architecture()[0] != "64bit",
         "Cross compile for windows can only be enabled on linux x86-64 platform",
     )
+    @unittest.skipIf(
+        not (check_cross_compile_trt_win_lib()),
+        "TRT windows lib for cross compile not found",
+    )
     @pytest.mark.unit
     def test_dynamo_cross_compile_for_windows_multiple_output(self):
         class Add(torch.nn.Module):