Precision Issue with Linear Reference

[tkreiman]

Python version

3.12.8

fairchem-core version

1.4.1.dev13+g2e49a4ef

pytorch version

2.4.0

cuda version

12.1

Operating system version

No response

Minimal example

from fairchem.core.datasets import LmdbDataset
import torch
from torch_geometric.data import Batch
d = LmdbDataset(cfg)
linear_reference = torch.load(path_to_lr)['energy']

obj = d[100]
b = Batch.from_data_list([obj])
y = b.y

dereferenced = linear_reference.dereference(y.unsqueeze(0), b)
dereferenced_double = linear_reference.dereference(y.unsqueeze(0).double(), b)

print("Float")
print(y, linear_reference(dereferenced, b))
print("Diff meV:", (y - linear_reference(dereferenced, b)).abs() * 1000)
print("Double")
print(y, linear_reference(dereferenced_double, b))
print("Diff meV:", (y - linear_reference(dereferenced_double, b)).abs() * 1000)

Current behavior

Differences up to 20 meV for the true and linear referenced energy when using float. Difference goes away with double precision.

Expected Behavior

Much smaller difference (if not 0) even when using float.

Relevant files to reproduce this bug

No response