diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index 2918bba947..9a7ea14cfb 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/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) diff --git a/deepmd/pt/model/descriptor/repformer_layer.py b/deepmd/pt/model/descriptor/repformer_layer.py index 92e2404469..5270c94112 100644 --- a/deepmd/pt/model/descriptor/repformer_layer.py +++ b/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) diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index ad4ead4d74..f237088a16 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/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 diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index e939a2541b..ffd645f2b9 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/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 diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index c028230e9b..c3174a2011 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/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 diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index e82bb23dac..4492a6c6b5 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/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 diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index 072457b48f..49dbdaf027 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/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 diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 437a464709..c140527f31 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/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 diff --git a/deepmd/pt/utils/neighbor_stat.py b/deepmd/pt/utils/neighbor_stat.py index d427dc758a..7d52bfaae1 100644 --- a/deepmd/pt/utils/neighbor_stat.py +++ b/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 diff --git a/source/tests/pt/model/test_jit.py b/source/tests/pt/model/test_jit.py index 248ccf9173..1f1034c330 100644 --- a/source/tests/pt/model/test_jit.py +++ b/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()