mixing creation and use of float32 and float64 calculators leads to a crash

[bernstei]

I suppose this might be too weird of an edge case to fix, but the following code

from ase.atoms import Atoms

from mace.calculators import mace_mp

atoms = Atoms('H', cell=[2] * 3, pbc=[True] * 3)

calc64 = mace_mp("small", device="cpu", default_dtype="float64")
calc32 = mace_mp("small", device="cpu")

calc32.get_descriptors(atoms=atoms)

calc64.get_descriptors(atoms=atoms)

results in the following output and error. I'm guessing that something is getting cached module-wide, and that ends up with mixed type in the error message

/home/cluster/bernstei/.local/lib/python3.11/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  _Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))
Using Materials Project MACE for MACECalculator with /home/cluster/bernstei/.cache/mace/20231210mace128L0_energy_epoch249model
Using float64 for MACECalculator, which is slower but more accurate. Recommended for geometry optimization.
/home/cluster/bernstei/.local/lib/python3.11/site-packages/mace/calculators/mace.py:139: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  torch.load(f=model_path, map_location=device)
Using Materials Project MACE for MACECalculator with /home/cluster/bernstei/.cache/mace/20231210mace128L0_energy_epoch249model
Using float32 for MACECalculator, which is faster but less accurate. Recommended for MD. Use float64 for geometry optimization.
/home/cluster/bernstei/.local/lib/python3.11/site-packages/mace/calculators/mace.py:139: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  torch.load(f=model_path, map_location=device)
Default dtype float32 does not match model dtype float64, converting models to float32.
Traceback (most recent call last):
  File "/home/cluster/bernstei/t.py", line 12, in <module>
    calc64.get_descriptors(atoms=atoms)
  File "/home/cluster/bernstei/.local/lib/python3.11/site-packages/mace/calculators/mace.py", line 412, in get_descriptors
    descriptors = [model(batch.to_dict())["node_feats"] for model in self.models]
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cluster/bernstei/.local/lib/python3.11/site-packages/mace/calculators/mace.py", line 412, in <listcomp>
    descriptors = [model(batch.to_dict())["node_feats"] for model in self.models]
                   ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/Software/python/conda/torch/2.5.1/cpu/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/Software/python/conda/torch/2.5.1/cpu/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cluster/bernstei/.local/lib/python3.11/site-packages/mace/modules/models.py", line 382, in forward
    node_e0 = self.atomic_energies_fn(data["node_attrs"])[
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/Software/python/conda/torch/2.5.1/cpu/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/Software/python/conda/torch/2.5.1/cpu/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cluster/bernstei/.local/lib/python3.11/site-packages/mace/modules/blocks.py", line 197, in forward
    return torch.matmul(x, torch.atleast_2d(self.atomic_energies).T)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: expected m1 and m2 to have the same dtype, but got: float != double

[ElliottKasoar]

I've encountered similar issues trying to compare results from MACE with other models e.g. CHGNet.

When the MACE calculators are initialised, they call torch.set_default_dtype(dtype) (torch.get_default_dtype() defaults to torch.float32, so calc64 would still error without this).

My work around so far has just been to call set_default_dtype again between calculations.

Assuming the default needs to be used, i.e. it's not possible to explicitly set all the relevant types, the only other solution I've considered is adding a check in the calculate function call, but it felt a bit unnecessary. Then again, calling get_default_dtype ought to be insignificant compared to running the model itself, so maybe it's not the worst idea.