Feat: add model format for dpa1

iProzd · Feb 1, 2024 · 51643af · 51643af
1 parent afb440a
commit 51643af
Show file tree

Hide file tree

Showing 16 changed files with 1,556 additions and 87 deletions.
diff --git a/deepmd/model_format/__init__.py b/deepmd/model_format/__init__.py
@@ -14,6 +14,8 @@
     EmbeddingNet,
     FittingNet,
     NativeLayer,
+    EmbdLayer,
+    LayerNorm,
     NativeNet,
     NetworkCollection,
     load_dp_model,
@@ -35,10 +37,14 @@
 from .se_e2_a import (
     DescrptSeA,
 )
+from .dpa1 import DescrptDPA1
 
 __all__ = [
     "InvarFitting",
     "DescrptSeA",
+    "DescrptDPA1",
+    "EmbdLayer",
+    "LayerNorm",
     "EnvMat",
     "make_multilayer_network",
     "make_embedding_network",