diff --git a/news/570.update.rst b/news/570.update.rst new file mode 100644 index 00000000..1121056a --- /dev/null +++ b/news/570.update.rst @@ -0,0 +1,2 @@ +Add automatic increase of recursion limit in the ``torch`` hook to ensure that +recursion limit is at least 5000 if ``torch`` 2.0.0 or later is detected. \ No newline at end of file diff --git a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py index 93888ab3..760e9f8f 100644 --- a/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py +++ b/src/_pyinstaller_hooks_contrib/hooks/stdhooks/hook-torch.py @@ -10,6 +10,15 @@ # SPDX-License-Identifier: GPL-2.0-or-later # ------------------------------------------------------------------ -from PyInstaller.utils.hooks import get_package_paths +from PyInstaller.utils.hooks import logger, get_package_paths, is_module_satisfies datas = [(get_package_paths('torch')[1],"torch"),] + +# With torch 2.0.0, PyInstaller's modulegraph analysis hits the recursion limit. +# So, unless the user has already done so, increase it automatically. +if is_module_satisfies('torch >= 2.0.0'): + import sys + new_limit = 5000 + if sys.getrecursionlimit() < new_limit: + logger.info("hook-torch: raising recursion limit to %d", new_limit) + sys.setrecursionlimit(new_limit)