adopt for SchNet

src/graph_pes/models/scaling.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import warnings
 from abc import ABC, abstractmethod
 from typing import Sequence
@@ -12,7 +14,7 @@
 from graph_pes.nn import PerElementParameter
 
 
-class ScaledPESModel(GraphPESModel, ABC):
+class UnScaledPESModel(GraphPESModel, ABC):
     """
     An abstract base class for all PES models implementations that are best
     suited to making raw predictions that with ~unit variance. By inheriting
@@ -66,11 +68,12 @@ def pre_fit(self, graphs: LabelledGraphDataset | Sequence[LabelledGraph]):
         # use Ridge regression to calculate standard deviations in the
         # per-element contributions to the total energy
         if "energy" in graph_batch:
-            _, variances = guess_per_element_mean_and_var(
-                graph_batch["energy"], graph_batch
-            )
-            for Z, var in variances.items():
-                self._per_element_scaling[Z] = var**0.5
+            with torch.no_grad():
+                _, variances = guess_per_element_mean_and_var(
+                    graph_batch["energy"], graph_batch
+                )
+                for Z, var in variances.items():
+                    self._per_element_scaling[Z] = var**0.5
 
         else:
             model_name = self.__class__.__name__

src/graph_pes/models/schnet.py

@@ -4,9 +4,9 @@
 from torch import Tensor, nn
 from torch_geometric.nn import MessagePassing
 
-from graph_pes.core import GraphPESModel
 from graph_pes.graphs import AtomicGraph
 from graph_pes.graphs.operations import neighbour_distances
+from graph_pes.models.scaling import UnScaledPESModel
 from graph_pes.nn import MLP, PerElementEmbedding, ShiftedSoftplus
 
 from .distances import DistanceExpansion, GaussianSmearing
@@ -167,7 +167,7 @@ def forward(
         return self.mlp(h)
 
 
-class SchNet(GraphPESModel):
+class SchNet(UnScaledPESModel):
     r"""
     The `SchNet <https://arxiv.org/abs/1706.08566>`_ model: a pairwise, scalar,
     message passing GNN.
@@ -237,7 +237,7 @@ def __init__(
             activation=ShiftedSoftplus(),
         )
 
-    def predict_local_energies(self, graph: AtomicGraph) -> torch.Tensor:
+    def predict_unscaled_energies(self, graph: AtomicGraph) -> torch.Tensor:
         h = self.chemical_embedding(graph["atomic_numbers"])
 
         for interaction in self.interactions: