fix(jax): calculate virial in call_lower#4304

wanghan-iapcm merged 2 commits into

deepmodeling:develdeepmodeling/deepmd-kit:develfrom

njzjz:jax-call-lower-virialnjzjz/deepmd-kit:jax-call-lower-virialCopy head branch name to clipboard