From 0fa72ef55e7941b9a055fe3a5575ff33f534e806 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 31 Oct 2024 20:36:47 +0000
Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/jax/utils/serialization.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py
index 102b588a9e..ec2de3060e 100644
--- a/deepmd/jax/utils/serialization.py
+++ b/deepmd/jax/utils/serialization.py
@@ -68,7 +68,9 @@ def call_lower_with_fixed_do_atomic_virial(
                 )
 
             return jax_export.export(jax.jit(call_lower_with_fixed_do_atomic_virial))(
-                jax.ShapeDtypeStruct((nf, nloc + nghost, 3), jnp.float64),  # extended_coord
+                jax.ShapeDtypeStruct(
+                    (nf, nloc + nghost, 3), jnp.float64
+                ),  # extended_coord
                 jax.ShapeDtypeStruct((nf, nloc + nghost), jnp.int32),  # extended_atype
                 jax.ShapeDtypeStruct((nf, nloc, model.get_nnei()), jnp.int64),  # nlist
                 jax.ShapeDtypeStruct((nf, nloc + nghost), jnp.int64),  # mapping