diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py
index e72c0d1e0f..37b78d39c9 100644
--- a/deepmd/fit/ener.py
+++ b/deepmd/fit/ener.py
@@ -406,7 +406,7 @@ def build (self,
                 # 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)
             else:
-                inputs_zero = tf.zeros((nframes, self.dim_descrpt * natoms[0]), dtype=self.fitting_precision)
+                inputs_zero = tf.zeros_like(inputs, dtype=self.fitting_precision)
         
 
         if bias_atom_e is not None :