@@ -90,21 +90,28 @@ def _find_lib(name: str, paths: List[str]) -> str:
90
90
91
91
def _register_with_torch () -> None :
92
92
trtorch_dir = os .path .dirname (__file__ )
93
- linked_file = (
94
- "/ torchtrt.dll" if sys .platform .startswith ("win" ) else "/lib/ libtorchtrt.so"
93
+ linked_file = os . path . join (
94
+ "lib" , ( " torchtrt.dll" if sys .platform .startswith ("win" ) else "libtorchtrt.so" )
95
95
)
96
- linked_file_runtime = (
97
- "/torchtrt_runtime.dll"
98
- if sys .platform .startswith ("win" )
99
- else "/lib/libtorchtrt_runtime.so"
96
+ linked_file_runtime = os .path .join (
97
+ "lib" ,
98
+ (
99
+ "torchtrt_runtime.dll"
100
+ if sys .platform .startswith ("win" )
101
+ else "libtorchtrt_runtime.so"
102
+ ),
100
103
)
101
- if os .path .isfile (trtorch_dir + linked_file ):
104
+ linked_file_full_path = os .path .join (trtorch_dir , linked_file )
105
+ linked_file_runtime_full_path = os .path .join (trtorch_dir , linked_file_runtime )
106
+
107
+ if os .path .isfile (linked_file_full_path ):
102
108
assert ENABLED_FEATURES .torchscript_frontend
103
109
assert ENABLED_FEATURES .torch_tensorrt_runtime
104
- torch .ops .load_library (trtorch_dir + linked_file )
105
- elif os .path .isfile (trtorch_dir + linked_file_runtime ):
110
+ torch .ops .load_library (linked_file_full_path )
111
+
112
+ elif os .path .isfile (linked_file_runtime_full_path ):
106
113
assert ENABLED_FEATURES .torch_tensorrt_runtime
107
- torch .ops .load_library (trtorch_dir + linked_file_runtime )
114
+ torch .ops .load_library (linked_file_runtime_full_path )
108
115
109
116
110
117
_register_with_torch ()
0 commit comments