Avoid importing torch in one-click-installer (#4064)

This commit is contained in:
jllllll 2023-09-24 20:16:59 -05:00 committed by GitHub
parent 55a685d999
commit c0fca23cb9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -53,7 +53,16 @@ def cpu_has_avx2():
def torch_version():
from torch import __version__ as torver
for sitedir in site.getsitepackages():
if "site-packages" in sitedir and conda_env_path in sitedir:
site_packages_path = sitedir
break
if site_packages_path:
torch_version_file = open(os.path.join(site_packages_path, 'torch', 'version.py')).read().splitlines()
torver = [line for line in torch_version_file if '__version__' in line][0].split('__version__ = ')[1].strip("'")
else:
from torch import __version__ as torver
return torver