Skip to content

Commit 1a578d1

Browse files
committed
Fix TORCHAO_SKIP_LOADING_SO_FILES behavior (#3189)
**Summary:** Today, if users want to force skip loading the .so files, they can pass in `TORCHAO_SKIP_LOADING_SO_FILES=1`. However, if they pass in `TORCHAO_SKIP_LOADING_SO_FILES=0` or `TORCHAO_SKIP_LOADING_SO_FILES=false`, we will still skip loading these files. This commit fixes this behavior by: 1. Renaming this env var to `TORCHAO_FORCE_SKIP_LOADING_SO_FILES` 2. Only accepting value of "1" for this env var **Test Plan:** ``` $ TORCHAO_FORCE_SKIP_LOADING_SO_FILES=1 python -c "import torchao" Skipping import of cpp extensions due to TORCHAO_FORCE_SKIP_LOADING_SO_FILES=1 \# No effect $ TORCHAO_FORCE_SKIP_LOADING_SO_FILES=0 python -c "import torchao" $ TORCHAO_FORCE_SKIP_LOADING_SO_FILES=False python -c "import torchao" $ TORCHAO_FORCE_SKIP_LOADING_SO_FILES=false python -c "import torchao" ```
1 parent 937e3a4 commit 1a578d1

File tree

1 file changed

+14
-6
lines changed

1 file changed

+14
-6
lines changed

torchao/__init__.py

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,12 @@ def _parse_version(version_string):
4545

4646

4747
skip_loading_so_files = False
48-
if bool(os.getenv("TORCHAO_SKIP_LOADING_SO_FILES", False)):
48+
force_skip_loading_so_files = (
49+
os.getenv("TORCHAO_FORCE_SKIP_LOADING_SO_FILES", "0") == "1"
50+
)
51+
if force_skip_loading_so_files:
4952
# user override
50-
# users can set env var TORCH_INCOMPATIBLE=1 to skip loading .so files
53+
# users can set env var TORCHAO_FORCE_SKIP_LOADING_SO_FILES=1 to skip loading .so files
5154
# this way, if they are using an incompatbile torch version, they can still use the API by setting the env var
5255
skip_loading_so_files = True
5356
# if torchao version has "+git", assume it's locally built and we don't know
@@ -80,10 +83,15 @@ def _parse_version(version_string):
8083

8184

8285
if skip_loading_so_files:
83-
logger.warning(
84-
f"Skipping import of cpp extensions due to incompatible torch version {torch.__version__} for torchao version {__version__} \
85-
Please see GitHub issue #2919 for more info"
86-
)
86+
if force_skip_loading_so_files:
87+
logger.warning(
88+
"Skipping import of cpp extensions due to TORCHAO_FORCE_SKIP_LOADING_SO_FILES=1"
89+
)
90+
else:
91+
logger.warning(
92+
f"Skipping import of cpp extensions due to incompatible torch version {torch.__version__} for torchao version {__version__} \
93+
Please see https://github.com/pytorch/ao/issues/2919 for more info"
94+
)
8795
else:
8896
try:
8997
from pathlib import Path

0 commit comments

Comments
 (0)