Floating Point Inconsistency Error

[RongkunZhou]

Hello, I'm trying to replicate the experiment detailed in the Colab on my own server. But I am encountering a floating point inconsistency error. I set the 'default_dtype' in 'tutorial.yaml' to 'float64'. Here are the details:

Environment:

NequIP: 0.5.5
LAMMPS: 27 Jun 2024
PyTorch: 1.11.0+cu113
CUDA: 11.3
Python: 3.8.10
libtorch: deps-1.11.0+cu113
GPU: RTX3080Ti

Error Message:

LAMMPS (27 Jun 2024)
  using 12 OpenMP thread(s) per MPI task
Reading data file ...
  orthogonal box = (0 0 0) to (10.862 10.862 10.862)
  1 by 1 by 1 MPI processor grid
  reading atoms ...
  64 atoms
  read_data CPU = 0.001 seconds
Allegro is using input precision f and output precision d
Allegro: Loading model from ../si-deployed.pth
Allegro: Freezing TorchScript model...
Type mapping:
Allegro type | Allegro name | LAMMPS type | LAMMPS name
0 | Si | 1 | Si

CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE

Your simulation uses code contributions which should be cited:
- Type Label Framework: https://doi.org/10.1021/acs.jpcb.3c08419
The log file lists these citations in BibTeX format.

CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE-CITE

Neighbor list info ...
  update: every = 1 steps, delay = 5 steps, check = yes
  max neighbors/atom: 2000, page size: 100000
  master list distance cutoff = 6
  ghost atom cutoff = 7
  binsize = 3, bins = 4 4 4
  2 neighbor lists, perpetual/occasional/extra = 1 1 0
  (1) pair allegro, perpetual
      attributes: full, newton on, ghost
      pair build: full/bin/ghost
      stencil: full/ghost/bin/3d
      bin: standard
  (2) compute rdf, occasional
      attributes: half, newton on
      pair build: half/bin/atomonly/newton
      stencil: half/bin/3d
      bin: standard
Setting up Verlet run ...
  Unit style    : metal
  Current step  : 0
  Time step     : 0.001
Exception: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/nequip/nn/_grad_output.py", line 32, in forward
        _6 = torch.append(wrt_tensors, data[k])
      func0 = self.func
      data0 = (func0).forward(data, )
               ~~~~~~~~~~~~~~ <--- HERE
      of = self.of
      _7 = [torch.sum(data0[of])]
  File "code/__torch__/nequip/nn/_graph_mixin.py", line 28, in AD_sum_backward
    input1 = (radial_basis).forward(input0, )
    input2 = (spharm).forward(input1, )
    input3 = (allegro).forward(input2, )
              ~~~~~~~~~~~~~~~~ <--- HERE
    input4 = (edge_eng).forward(input3, )
    input5 = (edge_eng_sum).forward(input4, )
  File "code/__torch__/allegro/nn/_allegro.py", line 111, in mean_0
    _18 = torch.unsqueeze(torch.index(cutoff_coeffs, _17), -1)
    new_latents0 = torch.mul(_18, new_latents)
    latents1 = torch.index_copy(latents, 0, active_edges0, new_latents0)
               ~~~~~~~~~~~~~~~~ <--- HERE
    _19 = annotate(List[Optional[Tensor]], [active_edges0])
    weights = (_01).forward(torch.index(latents1, _19), )

Traceback of TorchScript, original code (most recent call last):
  File "/root/miniconda3/lib/python3.8/site-packages/nequip/nn/_grad_output.py", line 85, in forward
            wrt_tensors.append(data[k])
        # run func
        data = self.func(data)
               ~~~~~~~~~ <--- HERE
        # Get grads
        grads = torch.autograd.grad(
  File "/root/miniconda3/lib/python3.8/site-packages/nequip/nn/_graph_mixin.py", line 356, in AD_sum_backward
    def forward(self, input: AtomicDataDict.Type) -> AtomicDataDict.Type:
        for module in self:
            input = module(input)
                    ~~~~~~ <--- HERE
        return input
  File "/root/miniconda3/lib/python3.8/site-packages/allegro/nn/_allegro.py", line 531, in mean_0
                # Normal (non-residual) update
                # index_copy replaces, unlike index_add
                latents = torch.index_copy(latents, 0, active_edges, new_latents)
                          ~~~~~~~~~~~~~~~~ <--- HERE
    
            # From the latents, compute the weights for active edges:
RuntimeError: index_copy_(): self and source expected to have the same dtype, but got (self) Float and (source) Double

The error indicates that there is an inconsistency between the data types (Float vs Double) during an operation within the TorchScript interpreter. Is this issue related to the versions of the libraries I am using? How can I resolve the floating point inconsistency?

Thank you very much!