diff --git a/tensorflow_quantum/core/ops/load_module.py b/tensorflow_quantum/core/ops/load_module.py index a98a27422..6fa350523 100644 --- a/tensorflow_quantum/core/ops/load_module.py +++ b/tensorflow_quantum/core/ops/load_module.py @@ -18,6 +18,7 @@ import types from distutils.sysconfig import get_python_lib +from tensorflow.python.framework import errors from tensorflow.python.framework import load_library from tensorflow.python.platform import resource_loader @@ -47,7 +48,7 @@ def _load(self): try: path = resource_loader.get_path_to_datafile(self._name) self._module = load_library.load_op_library(path) - except: + except (IOError, errors.NotFoundError): path = os.path.join(get_python_lib(), "tensorflow_quantum/core/ops", self._name) self._module = load_library.load_op_library(path)