fix se_e2_r prod_force_virial zero input

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Jul 27, 2024 · 4aa24ae · 4aa24ae
1 parent dd56fe9
commit 4aa24ae
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py
@@ -512,10 +512,11 @@ def prod_force_virial(
         """
         [net_deriv] = tf.gradients(atom_ener, self.descrpt_reshape)
         tf.summary.histogram("net_derivative", net_deriv)
+        nf = tf.shape(self.nlist)[0]
         net_deriv_reshape = tf.reshape(
             net_deriv,
             [
-                np.asarray(-1, dtype=np.int64),
+                nf,
                 natoms[0] * np.asarray(self.ndescrpt, dtype=np.int64),
             ],
         )