fix: UT

deepmodeling · Oct 31, 2024 · 06c9e11 · 06c9e11
1 parent 435a09e
commit 06c9e11
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/deepmd/dpmodel/model/dp_zbl_model.py b/deepmd/dpmodel/model/dp_zbl_model.py
@@ -20,17 +20,18 @@
     make_model,
 )
 
-DPZBLEnergyModel_ = make_model(DPZBLLinearEnergyAtomicModel)
+DPZBLModel_ = make_model(DPZBLLinearEnergyAtomicModel)
 
 
 @BaseModel.register("zbl")
-class DPZBLModel(DPZBLEnergyModel_):
+class DPZBLModel(DPZBLModel_):
+    model_type = "ener"
     def __init__(
         self,
         *args,
         **kwargs,
     ):
-        DPZBLEnergyModel_.__init__(self, *args, **kwargs)
+        super().__init__(*args, **kwargs)
 
 
 @classmethod