fix: precommit

deepmodeling · Mar 14, 2024 · 06528f4 · 06528f4
1 parent e5b6cf0
commit 06528f4
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py
@@ -563,9 +563,7 @@ def build(
             outs = tf.reshape(
                 final_layer,
                 [
-                    tf.shape(inputs)[0],
-                    natoms[2 + type_i],
-                    self.numb_dos,
+                    tf.shape(inputs)[0], natoms[0], self.numb_dos
                 ],
             )
         # add bias