diff --git a/tensorflow_quantum/core/ops/load_module.py b/tensorflow_quantum/core/ops/load_module.py index b5002ad84..8728a0e84 100644 --- a/tensorflow_quantum/core/ops/load_module.py +++ b/tensorflow_quantum/core/ops/load_module.py @@ -15,6 +15,7 @@ """Module to load python op libraries.""" import os +import sysconfig from distutils.sysconfig import get_python_lib from tensorflow.python.framework import load_library @@ -41,6 +42,9 @@ def load_module(name): path = resource_loader.get_path_to_datafile(name) return load_library.load_op_library(path) except: - path = os.path.join(get_python_lib(), "tensorflow_quantum/core/ops", - name) + path = os.path.join(get_python_lib(prefix="/usr/local"), + "tensorflow_quantum/core/ops", name) + if not os.path.exists(path): + path = path.replace("python3/", + f"python{sysconfig.get_python_version()}/") return load_library.load_op_library(path)