diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py index 0feec63316..10fddec5c0 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py @@ -1,3 +1,4 @@ +import importlib import logging from typing import Any @@ -101,6 +102,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # TODO: Update this function when quantization is added def is_impure(self, node: torch.fx.node.Node) -> bool: - if node.target in (torch.ops.tensorrt.quantize_op.default,): - return True + if importlib.util.find_spec("modelopt"): + if node.target in (torch.ops.tensorrt.quantize_op.default,): + return True return False