Error: Not registered 'ProdEnvMatA' when load DeepMD model in Python

[Zhangzhang-up]

Hi, I want to use the graph.pb from DeepMD for Monte Carlo simulation. Initially, there was an out-of-memory issue. After referring to the discussion in "Out Of Memory Error using Python interface #2145", I modified the code to load the model only at the beginning. The part of code is as follows:
import tensorflow as tf
model_path = "G:/Deep/compress.pb"
graph = tf.Graph()
sess = tf.compat.v1.Session(graph=graph)
with graph.as_default():
with tf.io.gfile.GFile(model_path, 'rb') as f:
graph_def = tf.compat.v1.GraphDef()
graph_def.ParseFromString(f.read())
tf.import_graph_def(graph_def, name='')
def calculate_dp_potential_energy(molecule_1_coords, molecule_2_coords, alpha=6, min_center_dist=2.0,
min_atom_dist=1.2):......

However, I am continuously receiving an error, as shown below:
Traceback (most recent call last):
File "G:\Deep\Mayer\1127.py", line 27, in
graph, session = load_model(model_path)
File "G:\Deep\Mayer\1127.py", line 22, in load_model
tf.import_graph_def(graph_def, name="")
File "D:\python\lib\site-packages\tensorflow\python\util\deprecation.py", line 576, in new_func
return func(*args, **kwargs)
File "D:\python\lib\site-packages\tensorflow\python\framework\importer.py", line 406, in import_graph_def
return _import_graph_def_internal(
File "D:\python\lib\site-packages\tensorflow\python\framework\importer.py", line 510, in _import_graph_def_internal
results = c_api.TF_GraphImportGraphDefWithResults(
tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered 'ProdEnvMatA' in binary running on DESKTOP-JJOJU04. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) tf.contrib.resampler should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.

Please let me know if you could help me resolve this issue. Thanks!