fix(pt): make int rcut safe after jit op

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -74,8 +74,8 @@ def __init__(
         super().__init__(type_map, **kwargs)
         super().init_out_stat()
         self.tab_file = tab_file
-        self.rcut = rcut
-        self.tab = self._set_pairtab(tab_file, rcut)
+        self.rcut = float(rcut)
+        self.tab = self._set_pairtab(tab_file, self.rcut)
 
         self.type_map = type_map
         self.ntypes = len(type_map)

deepmd/pt/model/descriptor/repformer_layer.py

@@ -605,8 +605,8 @@ def __init__(
     ):
         super().__init__()
         self.epsilon = 1e-4  # protection of 1./nnei
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.ntypes = ntypes
         sel = [sel] if isinstance(sel, int) else sel
         self.nnei = sum(sel)

deepmd/pt/model/descriptor/repformers.py

@@ -193,8 +193,8 @@ def __init__(
             Random seed for parameter initialization.
         """
         super().__init__()
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.ntypes = ntypes
         self.nlayers = nlayers
         sel = [sel] if isinstance(sel, int) else sel

deepmd/pt/model/descriptor/se_a.py

@@ -395,8 +395,8 @@ def __init__(
         - axis_neuron: Number of columns of the sub-matrix of the embedding matrix.
         """
         super().__init__()
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.neuron = neuron
         self.filter_neuron = self.neuron
         self.axis_neuron = axis_neuron

deepmd/pt/model/descriptor/se_atten.py

@@ -149,8 +149,8 @@ def __init__(
         """
         super().__init__()
         del type
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.neuron = neuron
         self.filter_neuron = self.neuron
         self.axis_neuron = axis_neuron

deepmd/pt/model/descriptor/se_r.py

@@ -74,8 +74,8 @@ def __init__(
         **kwargs,
     ):
         super().__init__()
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.neuron = neuron
         self.filter_neuron = self.neuron
         self.set_davg_zero = set_davg_zero

deepmd/pt/model/descriptor/se_t.py

@@ -446,8 +446,8 @@ def __init__(
             Random seed for initializing the network parameters.
         """
         super().__init__()
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.neuron = neuron
         self.filter_neuron = self.neuron
         self.set_davg_zero = set_davg_zero

deepmd/pt/model/descriptor/se_t_tebd.py

@@ -512,8 +512,8 @@ def __init__(
         seed: Optional[Union[int, list[int]]] = None,
     ):
         super().__init__()
-        self.rcut = rcut
-        self.rcut_smth = rcut_smth
+        self.rcut = float(rcut)
+        self.rcut_smth = float(rcut_smth)
         self.neuron = neuron
         self.filter_neuron = self.neuron
         self.tebd_dim = tebd_dim

deepmd/pt/utils/neighbor_stat.py

@@ -44,7 +44,7 @@ def __init__(
         mixed_types: bool,
     ) -> None:
         super().__init__()
-        self.rcut = rcut
+        self.rcut = float(rcut)
         self.ntypes = ntypes
         self.mixed_types = mixed_types
 

source/tests/pt/model/test_jit.py

@@ -144,5 +144,27 @@ def tearDown(self):
         JITTest.tearDown(self)
 
 
+class TestEnergyModelDPA2IntRcut(unittest.TestCase, JITTest):
+    def setUp(self):
+        input_json = str(Path(__file__).parent / "water/se_atten.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        data_file = [str(Path(__file__).parent / "water/data/data_0")]
+        self.config["training"]["training_data"]["systems"] = data_file
+        self.config["training"]["validation_data"]["systems"] = data_file
+        self.config["model"] = deepcopy(model_dpa2)
+        self.config["model"]["descriptor"]["repinit"]["rcut"] = int(
+            self.config["model"]["descriptor"]["repinit"]["rcut"]
+        )
+        self.config["model"]["descriptor"]["repinit"]["rcut_smth"] = int(
+            self.config["model"]["descriptor"]["repinit"]["rcut_smth"]
+        )
+        self.config["training"]["numb_steps"] = 10
+        self.config["training"]["save_freq"] = 10
+
+    def tearDown(self):
+        JITTest.tearDown(self)
+
+
 if __name__ == "__main__":
     unittest.main()