diff --git a/python/tvm/relay/frontend/tensorflow_parser.py b/python/tvm/relay/frontend/tensorflow_parser.py index 771aed06ac10..22c6f94ae122 100644 --- a/python/tvm/relay/frontend/tensorflow_parser.py +++ b/python/tvm/relay/frontend/tensorflow_parser.py @@ -67,12 +67,16 @@ def _load_pb_file(self): def _get_tag_set(self): """Return the tag set of saved model, multiple metagraphs are not supported""" try: - from tensorflow.contrib.saved_model.python.saved_model import reader + from tensorflow.contrib.saved_model.python.saved_model.reader \ + import get_saved_model_tag_sets except ImportError: - raise ImportError( - "InputConfiguration: Unable to import saved_model.reader which is " - "required to get tag set from saved model.") - tag_sets = reader.get_saved_model_tag_sets(self._model_dir) + try: + from tensorflow.python.tools.saved_model_utils import get_saved_model_tag_sets + except ImportError: + raise ImportError( + "InputConfiguration: Unable to import get_saved_model_tag_sets which is " + "required to get tag set from saved model.") + tag_sets = get_saved_model_tag_sets(self._model_dir) return tag_sets[0] def _get_output_names(self):