Update tensorrt import exception.

1 files changed, 4 insertions, 3 deletions

diff --git a/tensorflow/contrib/tensorrt/__init__.py b/tensorflow/contrib/tensorrt/__init__.py
index d53a05827a..a07b297900 100644
--- a/tensorflow/contrib/tensorrt/__init__.py
+++ b/tensorflow/contrib/tensorrt/__init__.py
@@ -18,16 +18,17 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from tensorflow.python.framework import errors
+
 # pylint: disable=unused-import,wildcard-import,g-import-not-at-top
 try:
   from tensorflow.contrib.tensorrt.python import *
-except Exception as e:
+except errors.NotFoundError as e:
   no_trt_message = (
       '**** Failed to initialize TensorRT. This is either because the TensorRT'
       ' installation path is not in LD_LIBRARY_PATH, or because you do not have'
       ' it installed. If not installed, please go to'
       ' https://developer.nvidia.com/tensorrt to download and install'
       ' TensorRT ****')
-  print(no_trt_message)
-  raise e
+  raise e(no_trt_message)
 # pylint: enable=unused-import,wildcard-import,g-import-not-at-top