fix se_e2_r prod_force_virial zero input

njzjz · njzjz · commit 4aa24ae3fe80 · 2024-07-27T17:08:22.000-04:00
Signed-off-by: Jinzhe Zeng &lt;jinzhe.zeng@rutgers.edu&gt;
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),
             ],
         )

Original file line number	Diff line number	Diff line change
`@@ -512,10 +512,11 @@ def prod_force_virial(`
`512`	`512`	`"""`
`513`	`513`	`[net_deriv] = tf.gradients(atom_ener, self.descrpt_reshape)`
`514`	`514`	`tf.summary.histogram("net_derivative", net_deriv)`
	`515`	`+ nf = tf.shape(self.nlist)[0]`
`515`	`516`	`net_deriv_reshape = tf.reshape(`
`516`	`517`	`net_deriv,`
`517`	`518`	`[`
`518`		`- np.asarray(-1, dtype=np.int64),`
	`519`	`+ nf,`
`519`	`520`	`natoms[0] * np.asarray(self.ndescrpt, dtype=np.int64),`
`520`	`521`	`],`
`521`	`522`	`)`