diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py
index dca4e0d353..977900e3f8 100644
--- a/deepmd/fit/ener.py
+++ b/deepmd/fit/ener.py
@@ -332,6 +332,7 @@ def _build_lower(
     def build (self, 
                inputs : tf.Tensor,
                natoms : tf.Tensor,
+               nframes : tf.Tensor,
                input_dict : dict = {},
                reuse : bool = None,
                suffix : str = '', 
@@ -352,6 +353,8 @@ def build (self,
                 natoms[0]: number of local atoms
                 natoms[1]: total number of atoms held by this processor
                 natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
+        nframes : tf.Tensor
+                The number of frames
         reuse
                 The weights in the networks should be reused when get the variable.
         suffix
@@ -398,10 +401,11 @@ def build (self,
                                                 trainable = False,
                                                 initializer = tf.constant_initializer(self.aparam_inv_std))
             
-        inputs = tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]])
+        inputs = tf.reshape(inputs, [nframes, self.dim_descrpt * natoms[0]])
         if len(self.atom_ener):
             # only for atom_ener
-            inputs_zero = tf.zeros_like(inputs, dtype=self.fitting_precision)
+            # like inputs, but we don't want to add a dependency on inputs
+            inputs_zero = tf.zeros((nframes, self.dim_descrpt * natoms[0]), dtype=self.fitting_precision)
         
 
         if bias_atom_e is not None :
diff --git a/deepmd/model/ener.py b/deepmd/model/ener.py
index c179f00f7b..60ba79f404 100644
--- a/deepmd/model/ener.py
+++ b/deepmd/model/ener.py
@@ -144,6 +144,7 @@ def build (self,
 
         coord = tf.reshape (coord_, [-1, natoms[1] * 3])
         atype = tf.reshape (atype_, [-1, natoms[1]])
+        nframes = tf.shape(coord)[0]
 
         # type embedding if any
         if self.typeebd is not None:
@@ -187,6 +188,7 @@ def build (self,
 
         atom_ener = self.fitting.build (dout, 
                                         natoms, 
+                                        nframes,
                                         input_dict, 
                                         reuse = reuse, 
                                         suffix = suffix)