fix: sort aparam in the Python API (#2666)

deepmodeling · Jul 11, 2023 · 3a5f7cf · 3a5f7cf
1 parent 4744fc8
commit 3a5f7cf
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py
@@ -466,6 +466,10 @@ def _prepare_feed_dict(
             efield = np.reshape(efield, [nframes, natoms, 3])
             efield = efield[:, imap, :]
             efield = np.reshape(efield, [nframes, natoms * 3])
+        if self.has_aparam:
+            aparam = np.reshape(aparam, [nframes, natoms, fdim])
+            aparam = aparam[:, imap, :]
+            aparam = np.reshape(aparam, [nframes, natoms * fdim])
 
         # make natoms_vec and default_mesh
         natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type)