fix(jax): calculate virial in call_lower

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Nov 4, 2024 · 71255cb · 71255cb
1 parent bfbe2ed
commit 71255cb
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/deepmd/jax/model/base_model.py b/deepmd/jax/model/base_model.py
@@ -152,4 +152,6 @@ def eval_ce(
                     avr, [0, def_ndim + 1, *range(1, def_ndim + 1), def_ndim + 2]
                 )
                 model_predict[kk_derv_c] = extended_virial
+                # [nf, *def, 9]
+                model_predict[kk_derv_c + "_redu"] = jnp.sum(extended_virial, axis=1)
     return model_predict