Explicitly put the created extension into sys.modules for PT2.9. (#158)

vbharadwaj-bk · web-flow · commit b60224a052e3 · 2025-09-02T12:17:15.000-07:00
* Explicitly put the created extension into sys.modules for PT2.9.

* Replaced importlib with sys.

* Removed unused import.

* No need for default parameter.
diff --git a/openequivariance/extlib/__init__.py b/openequivariance/extlib/__init__.py
@@ -120,6 +120,8 @@ def postprocess(kernel):
             extra_include_paths=include_dirs,
             extra_ldflags=extra_link_args,
         )
+        if "generic_module" not in sys.modules:
+            sys.modules["generic_module"] = generic_module
 
     if not TORCH_COMPILE:
         warnings.warn(

Original file line number	Diff line number	Diff line change
`@@ -120,6 +120,8 @@ def postprocess(kernel):`
`120`	`120`	`extra_include_paths=include_dirs,`
`121`	`121`	`extra_ldflags=extra_link_args,`
`122`	`122`	`)`
	`123`	`+ if "generic_module" not in sys.modules:`
	`124`	`+ sys.modules["generic_module"] = generic_module`
`123`	`125`
`124`	`126`	`if not TORCH_COMPILE:`
`125`	`127`	`warnings.warn(`