Update ener_model.py

Signed-off-by: Anchor Yu <[email protected]>
deepmodeling · Dec 15, 2024 · 7dead9c · 7dead9c
1 parent 19d0cd3
commit 7dead9c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py
@@ -23,7 +23,6 @@
 )
 
 DPEnergyModel_ = make_model(DPEnergyAtomicModel)
-DPEnergyModel_ = make_hessian_model(DPEnergyModel_)
 
 
 @BaseModel.register("ener")
@@ -40,6 +39,8 @@ def __init__(
         self._hessian_enabled = False
 
     def enable_hessian(self):
+        self.__class__ = make_hessian_model(type(self))
+        self.hess_fitting_def = super(type(self), self).atomic_output_def()
         self.requires_hessian("energy")
         self._hessian_enabled = True