Fixes to some MD desfaults and torchscript issues

src/schnetpack/atomistic/aggregation.py

@@ -18,7 +18,7 @@ class Aggregation(nn.Module):
     def __init__(self, keys: List[str], output_key: str = "y"):
         super(Aggregation, self).__init__()
 
-        self.keys = keys
+        self.keys: List[str] = list(keys)
         self.output_key = output_key
 
     def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:

src/schnetpack/atomistic/model.py

@@ -53,8 +53,6 @@ class AtomisticModel(pl.LightningModule):
 
     """
 
-    required_derivatives: List[str]
-
     def __init__(
         self,
         representation: nn.Module,

src/schnetpack/atomistic/nuclear_repulsion.py

@@ -85,7 +85,7 @@ def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
         # Get exponents and coefficients, normalize the latter
         exponents = a_ij[..., None] * F.softplus(self.exponents)[None, ...]
         coefficients = F.softplus(self.coefficients)[None, ...]
-        coefficients = F.normalize(coefficients, p=1, dim=1)
+        coefficients = F.normalize(coefficients, p=1.0, dim=1)
 
         screening = torch.sum(
             coefficients * torch.exp(-exponents * d_ij[:, None]), dim=1
@@ -102,7 +102,7 @@ def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
         # Compute ZBL energy
         y_zbl = snn.scatter_add(repulsion * screening, idx_i, dim_size=n_atoms)
         y_zbl = snn.scatter_add(y_zbl, idx_m, dim_size=n_molecules)
-        y_zbl = 0.5 * self.ke * y_zbl.squeeze(-1)
+        y_zbl = 0.5 * self.ke * y_zbl
 
         result = {self.output_key: y_zbl}
 

src/schnetpack/md/calculators/schnetpack_calculator.py

@@ -92,7 +92,7 @@ def _load_model(self, model_file: str) -> AtomisticModel:
         """
 
         log.info("Loading model from {:s}".format(model_file))
-        model = torch.load(model_file).to(torch.float32)
+        model = torch.load(model_file).to(torch.float64)
         model = model.eval()
 
         if self.stress_label is not None:
@@ -101,7 +101,7 @@ def _load_model(self, model_file: str) -> AtomisticModel:
 
         if self.script_model:
             log.info("Converting model to torch script...")
-            model = model.to_torchscript(None, "script")
+            model = model.to_torchscript(file_path=None, method="script")
 
         log.info("Deactivating inference mode for simulation...")
         self._deactivate_inference_mode(model)
@@ -116,7 +116,7 @@ def _deactivate_inference_mode(model: AtomisticModel) -> AtomisticModel:
                     log.info("Found `AddOffsets` postprocessing module...")
                     log.info(
                         "Constant offset of {:20.11f} will be removed...".format(
-                            pp.mean.detach().cpu().numpy()[0]
+                            pp.mean.detach().cpu().numpy()
                         )
                     )
         model.inference_mode = False

src/schnetpack/md/md_configs/calculator/spk.yaml

@@ -11,4 +11,4 @@ stress_label: null
 script_model: true
 
 defaults:
-  neighbor_list: ase
+  - neighbor_list: ase