Merge pull request #29 from wcwitt/torchscript_develop

Changes required for LAMMPS integration
ACEsuit · Oct 11, 2022 · 7db50e5 · 7db50e5
2 parents 0caeced + 72b815a
commit 7db50e5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/mace/calculators/mace.py b/mace/calculators/mace.py
@@ -30,7 +30,7 @@ def __init__(
         self.results = {}
 
         self.model = torch.load(f=model_path, map_location=device)
-        self.r_max = self.model.r_max
+        self.r_max = self.model.r_max.item()
         self.device = torch_tools.init_device(device)
         self.energy_units_to_eV = energy_units_to_eV
         self.length_units_to_A = length_units_to_A

diff --git a/mace/modules/models.py b/mace/modules/models.py
@@ -53,8 +53,8 @@ def __init__(
         gate: Optional[Callable],
     ):
         super().__init__()
-        self.r_max = r_max
-        self.atomic_numbers = atomic_numbers
+        self.register_buffer("r_max", torch.tensor(r_max, dtype=torch.float64))
+        self.register_buffer("atomic_numbers", torch.tensor(atomic_numbers, dtype=torch.int64))
         # Embedding
         node_attr_irreps = o3.Irreps([(num_elements, (0, 1))])
         node_feats_irreps = o3.Irreps([(hidden_irreps.count(o3.Irrep(0, 1)), (0, 1))])