fix: tf register

deepmodeling · Apr 15, 2024 · ff30173 · ff30173
1 parent 752132d
commit ff30173
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py
@@ -22,8 +22,6 @@
 )
 
 
-@StandardModel.register("dipole")
-@StandardModel.register("polar")
 class TensorModel(StandardModel):
     """Tensor model.
 
@@ -243,12 +241,12 @@ class WFCModel(TensorModel):
     def __init__(self, *args, **kwargs) -> None:
         TensorModel.__init__(self, "wfc", *args, **kwargs)
 
-
+@StandardModel.register("dipole")
 class DipoleModel(TensorModel):
     def __init__(self, *args, **kwargs) -> None:
         TensorModel.__init__(self, "dipole", *args, **kwargs)
 
-
+@StandardModel.register("polar")
 class PolarModel(TensorModel):
     def __init__(self, *args, **kwargs) -> None:
         TensorModel.__init__(self, "polar", *args, **kwargs)