address comment

deepmodeling · Feb 14, 2022 · 22be0f5 · 22be0f5
1 parent a731e59
commit 22be0f5
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py
@@ -738,7 +738,7 @@ def _filter_lower(
             # we can safely return the final xyz_scatter filled with zero directly
             return tf.cast(tf.fill((natom, 4, outputs_size[-1]), 0.), self.filter_precision)
           # natom x nei_type_i x out_size
-          xyz_scatter = tf.cast(tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1])), self.filter_precision)  
+          xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1]))  
           # When using tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]) below
           # [588 24] -> [588 6 4] correct
           # but if sel is zero

diff --git a/deepmd/utils/network.py b/deepmd/utils/network.py
@@ -79,15 +79,13 @@ def one_layer(inputs,
                 if use_timestep :
                     if mixed_prec is not None and not final_layer:
                        idt = tf.cast(idt, get_precision(mixed_prec['compute_prec']))
-                    return tf.reshape(activation_fn(hidden), [-1, outputs_size]) * idt
+                    hidden = tf.reshape(activation_fn(hidden), [-1, outputs_size]) * idt
                 else :
-                    return tf.reshape(activation_fn(hidden), [-1, outputs_size])                    
-        else:
-            if useBN:
-                None
-                # return self._batch_norm(hidden, name=name+'_normalization', reuse=reuse)
-            else:
-                return hidden
+                    hidden = tf.reshape(activation_fn(hidden), [-1, outputs_size])                    
+
+        if mixed_prec is not None:
+            hidden = tf.cast(hidden, get_precision(mixed_prec['output_prec']))
+        return hidden
 
 
 def embedding_net_rand_seed_shift(
@@ -237,6 +235,8 @@ def embedding_net(xx,
                 xx = tf.concat([xx,xx], 1) + hidden
         else:
             xx = hidden
+    if mixed_prec is not None:
+        xx = tf.cast(xx, get_precision(mixed_prec['output_prec']))
     return xx
 
 def variable_summaries(var: tf.Variable, name: str):