fix(pt): fix get_dim for DescrptDPA1Compat

deepmodeling · Jul 23, 2024 · 05842da · 05842da
1 parent a6ea2c1
commit 05842da
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py
@@ -2194,6 +2194,14 @@ def __init__(
         else:
             self.embd_input_dim = 1
 
+    def get_dim_out(self) -> int:
+        """Returns the output dimension of this descriptor."""
+        return (
+            super().get_dim_out() + self.tebd_dim
+            if self.concat_output_tebd
+            else super().get_dim_out()
+        )
+
     def build(
         self,
         coord_: tf.Tensor,