update __init__

deepmodeling · Dec 2, 2024 · 617a258 · 617a258
1 parent 984b242
commit 617a258
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/deepmd/pd/model/descriptor/__init__.py b/deepmd/pd/model/descriptor/__init__.py
@@ -22,15 +22,21 @@
     DescrptBlockSeA,
     DescrptSeA,
 )
+from .se_t_tebd import (
+    DescrptBlockSeTTebd,
+    DescrptSeTTebd,
+)
 
 __all__ = [
     "BaseDescriptor",
     "DescriptorBlock",
     "DescrptBlockRepformers",
     "DescrptBlockSeA",
     "DescrptBlockSeAtten",
+    "DescrptBlockSeTTebd",
     "DescrptDPA1",
     "DescrptDPA2",
     "DescrptSeA",
+    "DescrptSeTTebd",
     "prod_env_mat",
 ]