[BUG] PT: rcut cannot be int

[njzjz]

Bug summary

When rcut is an int, an error throws: Expected a value of type 'float' for argument 'rcut' but instead found type 'int'.

DeePMD-kit Version

b86165d

Backend and its version

PyTorch 2.1.2.post300

How did you download the software?

Built from source

Input Files, Running Commands, Error Log, etc.

  File "/home/jz748/codes/deepmd-kit/deepmd/pt/utils/neighbor_stat.py", line 137, in __init__
    self.op = torch.jit.script(op)
  File "/home/jz748/anaconda3/lib/python3.10/site-packages/torch/jit/_script.py", line 1324, in script
    return torch.jit._recursive.create_script_module(
  File "/home/jz748/anaconda3/lib/python3.10/site-packages/torch/jit/_recursive.py", line 559, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/home/jz748/anaconda3/lib/python3.10/site-packages/torch/jit/_recursive.py", line 636, in create_script_module_impl
    create_methods_and_properties_from_stubs(
  File "/home/jz748/anaconda3/lib/python3.10/site-packages/torch/jit/_recursive.py", line 469, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(
RuntimeError:

extend_coord_with_ghosts(Tensor coord, Tensor atype, Tensor? cell, float rcut, Tensor? cell_cpu=None) -> ((Tensor, Tensor, Tensor)):
Expected a value of type 'float' for argument 'rcut' but instead found type 'int'.
:
  File "/home/jz748/codes/deepmd-kit/deepmd/pt/utils/neighbor_stat.py", line 79
        nloc = coord.shape[1]
        coord = coord.view(nframes, nloc * 3)
        extend_coord, extend_atype, _ = extend_coord_with_ghosts(
                                        ~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
            coord, atype, cell, self.rcut
        )

Steps to Reproduce

cd examples/water/se_e2_a

In input.json, change "rcut": 6.00 to "rcut": 6.

dp --pt train input.json

Further Information, Files, and Links

No response

[njzjz]

As it looks like a bug from PyTorch, I submit to pytorch/pytorch#130682.