fix(pt): fix get_dim for DescrptDPA1Compat

deepmd/tf/descriptor/se_atten.py

@@ -765,7 +765,14 @@
             type_embedding=type_embedding,
             atype=atype,
         )
-        layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0], self.get_dim_out()])
+        layer = tf.reshape(
+            layer,
+            [
+                tf.shape(inputs)[0],
+                natoms[0],
+                self.filter_neuron[-1] * self.n_axis_neuron,
+            ],
+        )
         qmat = tf.reshape(
             qmat, [tf.shape(inputs)[0], natoms[0], self.get_dim_rot_mat_1() * 3]
         )
@@ -2194,6 +2201,14 @@
         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,

source/tests/consistent/descriptor/common.py

@@ -49,6 +49,8 @@ def build_tf_descriptor(self, obj, natoms, coords, atype, box, suffix):
             {},
             suffix=suffix,
         )
+        # ensure get_dim_out gives the correct shape
+        t_des = tf.reshape(t_des, [1, natoms[0], obj.get_dim_out()])
         return [t_des], {
             t_coord: coords,
             t_type: atype,