diff --git a/py/torch_tensorrt/__init__.py b/py/torch_tensorrt/__init__.py index 015a31f465..4eef3bab5d 100644 --- a/py/torch_tensorrt/__init__.py +++ b/py/torch_tensorrt/__init__.py @@ -4,6 +4,7 @@ import sys import platform import warnings +from packaging import version from torch_tensorrt._version import ( __version__, __cuda_version__, @@ -93,8 +94,10 @@ def _find_lib(name, paths): from torch_tensorrt._TRTModuleNext import TRTModuleNext from torch_tensorrt import fx -from torch_tensorrt import dynamo -from torch_tensorrt.dynamo import torch_compile + +if version.parse(torch.__version__) >= version.parse("2.dev"): + from torch_tensorrt import dynamo + from torch_tensorrt.dynamo import torch_compile def _register_with_torch():