diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py
index fe7572524f..aef134da92 100644
--- a/deepmd/tf/fit/dos.py
+++ b/deepmd/tf/fit/dos.py
@@ -562,9 +562,7 @@ def build(
 
             outs = tf.reshape(
                 final_layer,
-                [
-                    tf.shape(inputs)[0], natoms[0], self.numb_dos
-                ],
+                [tf.shape(inputs)[0], natoms[0], self.numb_dos],
             )
         # add bias
         # self.atom_ener_before = outs