fix: Fix meta kernel ops import issue for Python only builds (#3322)

pytorch · Dec 17, 2024 · 544c545 · 544c545
1 parent 827eca5
commit 544c545
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/py/torch_tensorrt/dynamo/runtime/__init__.py b/py/torch_tensorrt/dynamo/runtime/__init__.py
@@ -1,7 +1,10 @@
+import torch_tensorrt
 from torch_tensorrt.dynamo.runtime._PythonTorchTensorRTModule import (  # noqa: F401
     PythonTorchTensorRTModule,
 )
 from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import (  # noqa: F401
     TorchTensorRTModule,
 )
-from torch_tensorrt.dynamo.runtime.register_fake_class import *
+
+if torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime:
+    from torch_tensorrt.dynamo.runtime.meta_ops.register_meta_ops 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
@@ -593,13 +593,12 @@ def run(self):
         )
     ]
 
-    packages += [
-        "torch_tensorrt.ts",
-    ]
+    packages += ["torch_tensorrt.ts", "torch_tensorrt.dynamo.runtime.meta_ops"]
 
     package_dir.update(
         {
             "torch_tensorrt.ts": "py/torch_tensorrt/ts",
+            "torch_tensorrt.dynamo.runtime.meta_ops": "py/torch_tensorrt/dynamo/runtime/meta_ops",
         }
     )