add cast_precision decorator

deepmd/descriptor/se_a.py

@@ -7,6 +7,7 @@
 import numpy as np
 
 from deepmd.common import (
+    cast_precision,
     get_activation_func,
     get_precision,
 )
@@ -119,7 +120,7 @@ class DescrptSeA(paddle.nn.Layer):
     .. [1] Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018.
        End-to-end symmetry preserving inter-atomic potential energy model for finite and extended
        systems. In Proceedings of the 32nd International Conference on Neural Information Processing
-       Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441–4451.
+       Systems (NIPS'18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.
     """
 
     def __init__(
@@ -893,7 +894,7 @@ def _filter_lower(
         is_exclude=False,
     ):
         """Input env matrix, returns R.G."""
-        outputs_size = [1] + self.filter_neuron
+        outputs_size = [1, *self.filter_neuron]
         # cut-out inputs
         # with natom x (nei_type_i x 4)
         inputs_i = paddle.slice(
@@ -972,7 +973,7 @@ def _filter_lower(
                 transpose_x=True,
             )
 
-    # @cast_precision
+    @cast_precision
     def _filter(
         self,
         inputs: paddle.Tensor,
@@ -1021,14 +1022,12 @@ def _filter(
         nframes = 1
         # natom x (nei x 4)
         shape = inputs.shape
-        outputs_size = [1] + self.filter_neuron
+        outputs_size = [1, *self.filter_neuron]
         outputs_size_2 = self.n_axis_neuron  # 16
         all_excluded = all(
-            [
-                (type_input, type_i) in self.exclude_types  #  set()
-                for type_i in range(self.ntypes)
-            ]
-        )  # False
+            (type_input, type_i) in self.exclude_types  #  set()
+            for type_i in range(self.ntypes)
+        )
         if all_excluded:
             # all types are excluded so result and qmat should be zeros
             # we can safaly return a zero matrix...