diff options
| -rw-r--r-- | run_engine.py | 17 |
1 files changed, 13 insertions, 4 deletions
diff --git a/run_engine.py b/run_engine.py index 700494d..6ac2ec9 100644 --- a/run_engine.py +++ b/run_engine.py @@ -2,7 +2,11 @@ import os import sys from random import randint import numpy as np -import tensorrt + +if 'TensorRT-5.0' in os.getenv('LD_LIBRARY_PATH'): + import tensorrt.legacy +else: + import tensorrt try: from PIL import Image @@ -19,9 +23,14 @@ pip(3) install tensorrt[examples] exit(1) try: - import tensorrt as trt - from tensorrt.parsers import caffeparser - from tensorrt.parsers import onnxparser + if 'TensorRT-5.0' in os.getenv('LD_LIBRARY_PATH'): + import tensorrt.legacy as trt + from tensorrt.legacy.parsers import caffeparser + from tensorrt.legacy.parsers import onnxparser + else: + import tensorrt as trt + from tensorrt.parsers import caffeparser + from tensorrt.parsers import onnxparser except ImportError as err: sys.stderr.write("""ERROR: failed to import module ({}) Please make sure you have the TensorRT Library installed |