Skip to content

Commit 5f8c976

Browse files
fix bug in deepmd.infer.deep_pot.DeepPot (#2731)
- fix bug in checking t_efield op in the graph - delete arg atomic in `_prepare_feed_dict` (otherwise, efield cannot be imported correctly when using this method) --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 91633f0 commit 5f8c976

File tree

1 file changed

+6
-6
lines changed

1 file changed

+6
-6
lines changed

deepmd/infer/deep_pot.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -114,31 +114,32 @@ def __init__(
114114
operations = [op.name for op in self.graph.get_operations()]
115115
# check if the graph has these operations:
116116
# if yes add them
117-
if "t_efield" in operations:
118-
self._get_tensor("t_efield:0", "t_efield")
117+
118+
if ("%s/t_efield" % load_prefix) in operations:
119+
self.tensors.update({"t_efield": "t_efield:0"})
119120
self.has_efield = True
120121
else:
121122
log.debug("Could not get tensor 't_efield:0'")
122123
self.t_efield = None
123124
self.has_efield = False
124125

125-
if "load/t_fparam" in operations:
126+
if ("%s/t_fparam" % load_prefix) in operations:
126127
self.tensors.update({"t_fparam": "t_fparam:0"})
127128
self.has_fparam = True
128129
else:
129130
log.debug("Could not get tensor 't_fparam:0'")
130131
self.t_fparam = None
131132
self.has_fparam = False
132133

133-
if "load/t_aparam" in operations:
134+
if ("%s/t_aparam" % load_prefix) in operations:
134135
self.tensors.update({"t_aparam": "t_aparam:0"})
135136
self.has_aparam = True
136137
else:
137138
log.debug("Could not get tensor 't_aparam:0'")
138139
self.t_aparam = None
139140
self.has_aparam = False
140141

141-
if "load/spin_attr/ntypes_spin" in operations:
142+
if ("%s/spin_attr/ntypes_spin" % load_prefix) in operations:
142143
self.tensors.update({"t_ntypes_spin": "spin_attr/ntypes_spin:0"})
143144
self.has_spin = True
144145
else:
@@ -399,7 +400,6 @@ def _prepare_feed_dict(
399400
atom_types,
400401
fparam=None,
401402
aparam=None,
402-
atomic=False,
403403
efield=None,
404404
mixed_type=False,
405405
):

0 commit comments

Comments
 (0)