diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py
index 10240ceb56..122dfd7442 100644
--- 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)