Skip to content

Commit 52a3a9a

Browse files
authored
Use faster method to find the torch Python module's __path__ (#8611)
I noticed that the previous method took over 1 second on my M1 Mac Pro, so I found this method, which takes under 100 ms. I believe the root cause is that actually importing torch takes a little time.
1 parent f87940d commit 52a3a9a

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

build/Utils.cmake

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -357,7 +357,7 @@ function(add_torch_to_cmake_prefix_path)
357357
endif()
358358
execute_process(
359359
COMMAND "${PYTHON_EXECUTABLE}" -c
360-
"import torch as _; print(_.__path__[0], end='')"
360+
"import importlib.util; print(importlib.util.find_spec('torch').submodule_search_locations[0])"
361361
OUTPUT_VARIABLE _tmp_torch_path
362362
ERROR_VARIABLE _tmp_torch_path_error
363363
RESULT_VARIABLE _tmp_torch_path_result COMMAND_ECHO STDERR

0 commit comments

Comments
 (0)