fix head selection

ACEsuit · Aug 27, 2024 · 529b261 · 529b261
1 parent b0bbd5b
commit 529b261
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/mace/calculators/lammps_mace.py b/mace/calculators/lammps_mace.py
@@ -19,7 +19,7 @@ def __init__(self, model, **kwargs):
             torch.tensor(
                 self.model.heads.index(kwargs.get("head", self.model.heads[-1])),
                 dtype=torch.long,
-            ),
+            ).unsqueeze(0),
         )
 
         for param in self.model.parameters():

diff --git a/mace/cli/create_lammps_model.py b/mace/cli/create_lammps_model.py
@@ -19,7 +19,7 @@ def parse_args():
         type=str,
         nargs="?",
         help="Head of the model to be converted to LAMMPS",
-        default="default",
+        default=None,
     )
     return parser.parse_args()
 
@@ -30,7 +30,9 @@ def main():
     head = args.head
     model = torch.load(model_path)
     model = model.double().to("cpu")
-    lammps_model = LAMMPS_MACE(model, head=head)
+    lammps_model = (
+        LAMMPS_MACE(model, head=head) if head is not None else LAMMPS_MACE(model)
+    )
     lammps_model_compiled = jit.compile(lammps_model)
     lammps_model_compiled.save(model_path + "-lammps.pt")