fix: Fix meta kernel import in python builds

pytorch · Dec 13, 2024 · 986793d · 986793d
1 parent d47db76
commit 986793d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
diff --git a/py/torch_tensorrt/dynamo/runtime/__init__.py b/py/torch_tensorrt/dynamo/runtime/__init__.py
@@ -7,4 +7,4 @@
 )
 
 if torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime:
-    from torch_tensorrt.dynamo.runtime.register_fake_class import *
+    from torch_tensorrt.dynamo.runtime.meta_ops.register_meta_ops.py import *
diff --git a/...rrt/dynamo/runtime/register_fake_class.py → ...amo/runtime/meta_ops/register_meta_ops.py b/...rrt/dynamo/runtime/register_fake_class.py → ...amo/runtime/meta_ops/register_meta_ops.py
diff --git a/setup.py b/setup.py
@@ -444,6 +444,7 @@ def run(self):
     "torch_tensorrt.dynamo.lowering.passes",
     "torch_tensorrt.dynamo.partitioning",
     "torch_tensorrt.dynamo.runtime",
+    "torch_tensorrt.dynamo.runtime.meta_ops",
     "torch_tensorrt.dynamo.tools",
     "torch_tensorrt.fx",
     "torch_tensorrt.fx.converters",
@@ -484,6 +485,13 @@ def run(self):
     "torch_tensorrt.runtime": "py/torch_tensorrt/runtime",
 }
 
+if not PY_ONLY:
+    package_dir.update(
+        {
+            "torch_tensorrt.dynamo.runtime.meta_ops": "py/torch_tensorrt/dynamo/runtime/meta_ops"
+        }
+    )
+
 package_data = {}
 
 if not (PY_ONLY or NO_TS):