Import tensorrt if available to import_pb_to_tensorboard.py for displaying TensorRT ops

1 files changed, 7 insertions, 0 deletions

diff --git a/tensorflow/python/tools/import_pb_to_tensorboard.py b/tensorflow/python/tools/import_pb_to_tensorboard.py
index 00de044505..d1f9cd87b3 100755
--- a/tensorflow/python/tools/import_pb_to_tensorboard.py
+++ b/tensorflow/python/tools/import_pb_to_tensorboard.py
@@ -29,6 +29,13 @@ from tensorflow.python.platform import app
 from tensorflow.python.platform import gfile
 from tensorflow.python.summary import summary
 
+# Try importing TensorRT ops if available
+# pylint: disable=unused-import,trailing-whitespace
+try:
+  import tensorflow.contrib.tensorrt as trt 
+except ImportError:
+  pass
+# pylint: enable=unused-import,trailing-whitespace
 
 def import_to_tensorboard(model_dir, log_dir):
   """View an imported protobuf model (`.pb` file) as a graph in Tensorboard.