Skip to content

fix(torch_utils): explicitly inject cached build directory into sys.path - #661

Open
rahulsdhar wants to merge 1 commit into
NVlabs:mainfrom
rahulsdhar:fix/python312-cuda-import
Open

rahulsdhar wants to merge 1 commit into
NVlabs:mainfrom
rahulsdhar:fix/python312-cuda-import

Conversation

@rahulsdhar

Copy link
Copy Markdown

Description

This PR resolves a silent compilation blocker where importlib.import_module throws a ModuleNotFoundError: No module named 'bias_act_plugin' (and subsequently upfirdn2d_plugin) when trying to generate images on modern Python setups (specifically Python 3.11 and 3.12).

Root Cause

While torch.utils.cpp_extension.load() successfully compiles the CUDA/C++ source code into the .cache/torch_extensions/... hashed directory, modern Python environments do not always inherit or automatically discover this dynamic path via standard PyTorch path registration hooks.

Solution

Explicitly check if cached_build_dir is evaluated, and if so, safely insert it into sys.path right before invoking importlib.import_module(module_name). This mirrors standard dynamic linking paths and guarantees cross-compatibility with newer Python interpreters.

Testing

Verified on:

  • OS: Ubuntu / Linux
  • Python Version: 3.12
  • Command Run: python3 gen_images.py --outdir=out --trunc=1 --seeds=2 --network=...
  • Result: Successfully compiled, loaded modules, and generated outputs without breaking backward compatibility.

Fixes a ModuleNotFoundError when compiling custom CUDA extensions on modern Python environments (e.g., Python 3.11/3.12).

While torch.utils.cpp_extension.load compiles the custom plugins successfully into the hashed caching directory, modern PyTorch/Python path hooks sometimes fail to register this dynamic path automatically. Adding the verified cached_build_dir directly to sys.path guarantees that importlib can find and resolve the plugin seamlessly.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant