[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

deepmd/pt/model/descriptor/repformers.py

@@ -28,6 +28,9 @@
 from deepmd.pt.utils.exclude_mask import (
     PairExcludeMask,
 )
+from deepmd.pt.utils.spin import (
+    concat_switch_virtual,
+)
 from deepmd.pt.utils.utils import (
     ActivationFn,
 )
@@ -41,9 +44,6 @@
 from .repformer_layer import (
     RepformerLayer,
 )
-from deepmd.pt.utils.spin import (
-    concat_switch_virtual,
-)
 
 if not hasattr(torch.ops.deepmd, "border_op"):
 
@@ -493,7 +493,9 @@ def forward(
                 g1_ext = ret[0].unsqueeze(0)
                 if has_spin:
                     g1_real_ext, g1_virtual_ext = torch.split(g1_ext, [ng1, ng1], dim=2)
-                    g1_ext = concat_switch_virtual(g1_real_ext, g1_virtual_ext, real_nloc)
+                    g1_ext = concat_switch_virtual(
+                        g1_real_ext, g1_virtual_ext, real_nloc
+                    )
             g1, g2, h2 = ll.forward(
                 g1_ext,
                 g2,

deepmd/pt/model/model/spin_model.py

@@ -15,6 +15,9 @@
 from deepmd.pt.model.atomic_model import (
     DPAtomicModel,
 )
+from deepmd.pt.utils.spin import (
+    concat_switch_virtual,
+)
 from deepmd.pt.utils.utils import (
     to_torch_tensor,
 )
@@ -24,9 +27,6 @@
 from deepmd.utils.spin import (
     Spin,
 )
-from deepmd.pt.utils.spin import (
-    concat_switch_virtual,
-)
 
 from .make_model import (
     make_model,

deepmd/pt/utils/spin.py

@@ -1,4 +1,8 @@
-from typing import Optional
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    Optional,
+)
+
 import torch
 
 
@@ -24,12 +28,8 @@ def concat_switch_virtual(
         device=extended_tensor.device,
     )
     extended_tensor_updated[:, :nloc] = extended_tensor[:, :nloc]
-    extended_tensor_updated[:, nloc : nloc + nloc] = extended_tensor_virtual[
-        :, :nloc
-    ]
-    extended_tensor_updated[:, nloc + nloc : nloc + nall] = extended_tensor[
-        :, nloc:
-    ]
+    extended_tensor_updated[:, nloc : nloc + nloc] = extended_tensor_virtual[:, :nloc]
+    extended_tensor_updated[:, nloc + nloc : nloc + nall] = extended_tensor[:, nloc:]
     extended_tensor_updated[:, nloc + nall :] = extended_tensor_virtual[:, nloc:]
     # nloc + nloc + nghost + nghost
     if recv_num is not None:
@@ -59,4 +59,4 @@ def concat_switch_virtual(
             ] = extended_tensor_virtual[
                 :, nloc + origin_prefix_sum[i] : nloc + origin_prefix_sum[i + 1]
             ]
-    return extended_tensor_updated.view(out_shape)
+    return extended_tensor_updated.view(out_shape)