Make k_Z a torch parameter and fix module methods.

emle/models/_ani.py

@@ -144,6 +144,7 @@ def __init__(
                 raise TypeError("'device' must be of type 'torch.device'")
         else:
             device = _torch.get_default_device()
+        self._device = device
 
         if dtype is not None:
             if not isinstance(dtype, _torch.dtype):
@@ -229,8 +230,15 @@ def __init__(
                 except:
                     pass
 
+        # Add a hook to the ANI2x model to capture the AEV features.
+        self._add_hook()
+
+    def _add_hook(self):
+        """
+        Add a hook to the ANI2x model to capture the AEV features.
+        """
         # Assign a tensor attribute that can be used for assigning the AEVs.
-        self._ani2x.aev_computer._aev = _torch.empty(0, device=device)
+        self._ani2x.aev_computer._aev = _torch.empty(0, device=self._device)
 
         # Hook the forward pass of the ANI2x model to get the AEV features.
         # Note that this currently requires a patched versions of TorchANI and NNPOps.
@@ -261,6 +269,13 @@ def to(self, *args, **kwargs):
         """
         self._emle = self._emle.to(*args, **kwargs)
         self._ani2x = self._ani2x.to(*args, **kwargs)
+
+        # Check for a device type in args and update the device attribute.
+        for arg in args:
+            if isinstance(arg, _torch.device):
+                self._device = arg
+                break
+
         return self
 
     def cpu(self, **kwargs):
@@ -269,6 +284,7 @@ def cpu(self, **kwargs):
         """
         self._emle = self._emle.cpu(**kwargs)
         self._ani2x = self._ani2x.cpu(**kwargs)
+        self._device = _torch.device("cpu")
         return self
 
     def cuda(self, **kwargs):
@@ -277,6 +293,7 @@ def cuda(self, **kwargs):
         """
         self._emle = self._emle.cuda(**kwargs)
         self._ani2x = self._ani2x.cuda(**kwargs)
+        self._device = _torch.device("cuda")
         return self
 
     def double(self):
@@ -306,6 +323,9 @@ def float(self):
             except:
                 pass
 
+        # Re-append the hook.
+        self._add_hook()
+
         return self
 
     def forward(self, atomic_numbers, charges_mm, xyz_qm, xyz_mm):

emle/models/_emle.py

@@ -353,17 +353,6 @@ def __init__(
             dtype=dtype,
         )
 
-    def _to_dict(self):
-        """
-        Return the configuration of the module as a dictionary.
-        """
-        return {
-            "model": self._model,
-            "method": self._method,
-            "species": self._species_map.tolist(),
-            "alpha_mode": self._alpha_mode,
-        }
-
     def to(self, *args, **kwargs):
         """
         Performs Tensor dtype and/or device conversion on the model.
@@ -393,7 +382,7 @@ def cuda(self, **kwargs):
         self._emle_base = self._emle_base.cuda(**kwargs)
 
         # Update the device attribute.
-        self._device = self._species_map.device
+        self._device = self._q_total.device
 
         return self
 
@@ -408,7 +397,7 @@ def cpu(self, **kwargs):
         self._emle_base = self._emle_base.cpu()
 
         # Update the device attribute.
-        self._device = self._species_map.device
+        self._device = self._q_total.device
 
         return self
 

emle/models/_emle_base.py

@@ -207,7 +207,7 @@ def __init__(
         self.a_Thole = _torch.nn.Parameter(params["a_Thole"])
         self.ref_values_s = _torch.nn.Parameter(params["ref_values_s"])
         self.ref_values_chi = _torch.nn.Parameter(params["ref_values_chi"])
-        k_Z = _torch.nn.Parameter(params["k_Z"])
+        self.k_Z = _torch.nn.Parameter(params["k_Z"])
 
         if self._alpha_mode == "reference":
             try:
@@ -270,7 +270,6 @@ def __init__(
         self.register_buffer("_c_s", c_s)
         self.register_buffer("_c_chi", c_chi)
         self.register_buffer("_c_sqrtk", c_sqrtk)
-        self.register_buffer("_k_Z", k_Z)
 
         # Initalise an empty AEV tensor to use to store the AEVs in parent models.
         # If AEVs are computed externally, then this tensor will be set by the
@@ -292,7 +291,6 @@ def to(self, *args, **kwargs):
         self._c_s = self._c_s.to(*args, **kwargs)
         self._c_chi = self._c_chi.to(*args, **kwargs)
         self._c_sqrtk = self._c_sqrtk.to(*args, **kwargs)
-        self._k_Z = self._k_Z.to(*args, **kwargs)
 
         # Check for a device type in args and update the device attribute.
         for arg in args:
@@ -320,7 +318,7 @@ def cuda(self, **kwargs):
         self._c_s = self._c_s.cuda(**kwargs)
         self._c_chi = self._c_chi.cuda(**kwargs)
         self._c_sqrtk = self._c_sqrtk.cuda(**kwargs)
-        self._k_Z = self._k_Z.cuda(**kwargs)
+        self.k_Z = self.k_Z.cuda(**kwargs)
 
         # Update the device attribute.
         self._device = self._species_map.device
@@ -345,7 +343,6 @@ def cpu(self, **kwargs):
         self._c_s = self._c_s.cpu(**kwargs)
         self._c_chi = self._c_chi.cpu(**kwargs)
         self._c_sqrtk = self._c_sqrtk.cpu(**kwargs)
-        self._k_Z = self._k_Z.cpu(**kwargs)
 
         # Update the device attribute.
         self._device = self._species_map.device
@@ -367,7 +364,7 @@ def double(self):
         self._c_s = self._c_s.double()
         self._c_chi = self._c_chi.double()
         self._c_sqrtk = self._c_sqrtk.double()
-        self._k_Z = self._k_Z.double()
+        self.k_Z = _torch.nn.Parameter(self.k_Z.double())
         return self
 
     def float(self):
@@ -385,7 +382,7 @@ def float(self):
         self._c_s = self._c_s.float()
         self._c_chi = self._c_chi.float()
         self._c_sqrtk = self._c_sqrtk.float()
-        self._k_Z = self._k_Z.float()
+        self.k_Z = _torch.nn.Parameter(self.k_Z.float())
         return self
 
     def forward(self, atomic_numbers, xyz_qm, q_total):
@@ -445,7 +442,7 @@ def forward(self, atomic_numbers, xyz_qm, q_total):
         q = self._get_q(r_data, s, chi, q_total, mask)
         q_val = q - q_core
 
-        k = self._k_Z[species_id]
+        k = self.k_Z[species_id]
 
         if self._alpha_mode == "reference":
             k_scale = (