Skip to content

Commit 01232b9

Browse files
authored
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 ed4cd34 commit 01232b9

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
@@ -83,10 +86,15 @@ def _parse_version(version_string):
8386

8487

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

0 commit comments

Comments
 (0)