apply flax_module to energy model

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Oct 21, 2024 · 752d1e7 · 752d1e7
1 parent fa9a000
commit 752d1e7
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/deepmd/jax/model/ener_model.py b/deepmd/jax/model/ener_model.py
@@ -7,12 +7,16 @@
 from deepmd.jax.atomic_model.dp_atomic_model import (
     DPAtomicModel,
 )
+from deepmd.jax.common import (
+    flax_module,
+)
 from deepmd.jax.model.base_model import (
     BaseModel,
 )
 
 
 @BaseModel.register("ener")
+@flax_module
 class EnergyModel(EnergyModelDP):
     def __setattr__(self, name: str, value: Any) -> None:
         if name == "atomic_model":