Merge pull request #30 from kzinovjev/emle-train

Provide species list in model file
chemle · Oct 23, 2024 · 7e54719 · 7e54719
2 parents 760ac56 + e2f2dbd
commit 7e54719
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/emle/models/_emle.py b/emle/models/_emle.py
@@ -310,11 +310,7 @@ def __init__(
             "ref_values_chi": _torch.tensor(
                 params["chi_ref"], dtype=dtype, device=device
             ),
-            "k_Z": (
-                _torch.tensor(params["k_Z"], dtype=dtype, device=device)
-                if "k_Z" in params
-                else None
-            ),
+            "k_Z": _torch.tensor(params["k_Z"], dtype=dtype, device=device),
             "sqrtk_ref": (
                 _torch.tensor(params["sqrtk_ref"], dtype=dtype, device=device)
                 if "sqrtk_ref" in params
@@ -348,7 +344,7 @@ def __init__(
             aev_computer=self._aev_computer,
             aev_mask=aev_mask,
             alpha_mode=self._alpha_mode,
-            species=self._species,
+            species=params.get("species", self._species),
             device=device,
             dtype=dtype,
         )