diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index a8e17255b..bafb9cfaa 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -105,6 +105,7 @@ jobs:
 
   docs:
     runs-on: ubuntu-latest
+    timeout-minutes: 20
     steps:
       - uses: actions/checkout@v5
       - uses: actions/setup-python@v5
@@ -112,9 +113,8 @@ jobs:
           python-version: "3.12"
           cache: pip
           cache-dependency-path: pyproject.toml
-      # nbsphinx shells out to pandoc for markdown cells; the Read the Docs
-      # image ships it, GitHub's runner image does not guarantee it.
-      - run: sudo apt-get update && sudo apt-get install -y pandoc
+      - uses: pandoc/actions/setup@v1
+      - run: pip install torch --index-url https://download.pytorch.org/whl/cpu
       - run: pip install -e ".[doc]"
       - if: github.event_name == 'schedule'
         run: pip install -e ".[mace]"
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index 29dbd7658..bdaab28a4 100644
--- a/.github/workflows/python-publish.yml
+++ b/.github/workflows/python-publish.yml
@@ -11,9 +11,6 @@ name: Upload Python Package
 on:
   release:
     types: [published]
-  # Dry run of the same job against TestPyPI: `git push origin testpypi-<n>`.
-  push:
-    tags: ["testpypi-*"]
 
 permissions:
   contents: read
@@ -33,18 +30,10 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install build
-    # TestPyPI accepts every version only once, so each test run uploads its own
-    # dev version and the release version stays free.
-    - name: Use a dev version for TestPyPI
-      if: github.event_name == 'push'
-      run: |
-        sed -i -E 's/^__version__ = "([^"]+)"/__version__ = "\1.dev${{ github.run_number }}"/' src/schnetpack/__init__.py
-        grep '^__version__' src/schnetpack/__init__.py
     - name: Build package
       run: python -m build
     - name: Publish package
       uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
       with:
         user: __token__
-        password: ${{ github.event_name == 'release' && secrets.PYPI_API_TOKEN || secrets.TEST_PYPI_API_TOKEN }}
-        repository_url: ${{ github.event_name == 'release' && 'https://upload.pypi.org/legacy/' || 'https://test.pypi.org/legacy/' }}
+        password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index ed57fe56d..f462d2b24 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -2,8 +2,25 @@
 ci:
   autoupdate_schedule: monthly
   autofix_prs: true
+  skip: [jupytext-sync]
 
 repos:
+  # The tutorials and how-tos are tracked as jupytext .py files; the paired .ipynb is
+  # gitignored (see examples/jupytext.toml). This hook makes sure that jupytext .py files
+  # are synced with locally modified notebooks.
+  - repo: local
+    hooks:
+      - id: jupytext-sync
+        name: jupytext sync
+        entry: >-
+          bash -c 'find examples/howtos examples/tutorials_deprecated -name "*.ipynb"
+          -not -path "*/.ipynb_checkpoints/*" -print0 | xargs -0 -r -n1 jupytext --sync'
+        language: python
+        additional_dependencies: ["jupytext==1.19.5"]
+        # the .ipynb is gitignored, so pre-commit never sees it staged
+        always_run: true
+        pass_filenames: false
+
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.16.2
     hooks:
diff --git a/docs/userguide/md.rst b/docs/userguide/md.rst
index c18777ec7..7afc1ac58 100644
--- a/docs/userguide/md.rst
+++ b/docs/userguide/md.rst
@@ -19,6 +19,11 @@ The inputs which need to be provided for every ``spkmd`` run are:
 * the path to a trained ML model (``calculator.model_file``)
 * and the cutoff used in the neighbor list (``calculator.neighbor_list.cutoff``)
 
+.. warning::
+
+   Model files, checkpoints and restart files are pickled PyTorch objects. Loading them executes code contained in
+   the file, so only use files from sources you trust.
+
 Assuming the model and structure file are present in the local directory, the  command line call would be::
 
     spkmd simulation_dir=mdtut_cli system.molecule_file=md_ethanol.xyz calculator.model_file=md_ethanol.model calculator.neighbor_list.cutoff=5.0
diff --git a/examples/README.md b/examples/README.md
index e5ca8061b..8fe0905c6 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -14,15 +14,18 @@ jupytext --to notebook howtos/howto_batchwise_relaxations.py
 ```
 
 This writes `howto_batchwise_relaxations.ipynb` next to it, which you open in Jupyter as
-usual. The generated notebook is git-ignored, so any outputs you produce stay local.
-
-With jupytext installed you can also skip the conversion and open the `.py` directly from
-the JupyterLab file browser (*right-click → Open With → Notebook*). Your edits are then
+usual. With jupytext installed you can also skip the conversion and open the `.py` directly from the JupyterLab file browser (*right-click → Open With → Notebook*). Your edits are then
 written back to the `.py` automatically.
 
 If you just want to read them, the rendered versions are at
 [schnetpack.readthedocs.io](https://schnetpack.readthedocs.io).
 
+**For developers:** You can also edit the `.ipynb` in an editor that does not know about jupytext, such as
+VS Code. The `jupytext-sync` pre-commit hook writes every paired notebook back into its
+`.py` before a commit. If it rewrote a `.py`, the commit fails: `git add` the `.py` and
+commit again. A notebook you created from scratch gets a new `.py` that you have to
+`git add` yourself.
+
 ## Which ones CI executes
 
 Read the Docs and pull requests only render the examples. The weekly `docs` CI job
diff --git a/examples/benchmarks/batchwise_relaxation/plot_benchmark.py b/examples/benchmarks/batchwise_relaxation/plot_benchmark.py
deleted file mode 100644
index 8a10c8dd0..000000000
--- a/examples/benchmarks/batchwise_relaxation/plot_benchmark.py
+++ /dev/null
@@ -1,354 +0,0 @@
-"""Plot the relaxation time against the batch size, from the output of ``run_benchmark.py``.
-
-Prints the mean wall time of both optimizers per batch size with its confidence
-interval and the speedup, and plots both curves on log-log axes. With only a handful of
-repeats the interval is a Student-t interval, which is noticeably wider than the
-standard error: increase ``--repeats`` of the benchmark run to tighten it.
-"""
-
-import argparse
-import json
-import os
-
-import matplotlib.pyplot as plt
-import numpy as np
-from scipy import stats
-
-METHODS = ("batch-wise", "sequential")
-
-SURFACE = "#fcfcfb"
-INK = "#0b0b0b"
-INK_SECONDARY = "#52514e"
-MUTED = "#898781"
-GRID = "#e1e0d9"
-SERIES = {"batch-wise": "#2a78d6", "sequential": "#eb6834"}
-
-# "screen" to inspect a run, "paper" for a single column of a two-column journal: no
-# title (the caption carries it), print font sizes, every other batch size labelled
-LAYOUTS = {
-    "screen": dict(
-        figsize=(7.2, 4.8),
-        dpi=130,
-        surface=SURFACE,
-        font=10,
-        tick_font=9,
-        line=2,
-        marker=6,
-        pad=8,
-        tick_every=1,
-        title=True,
-    ),
-    "paper": dict(
-        figsize=(3.37, 2.5),
-        dpi=300,
-        surface="white",
-        font=8,
-        tick_font=7,
-        line=1.3,
-        marker=3.5,
-        pad=4,
-        tick_every=2,
-        title=False,
-    ),
-}
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
-    parser.add_argument(
-        "--input",
-        default="./batchwise_relaxation_benchmark",
-        help="output directory of run_benchmark.py",
-    )
-    parser.add_argument(
-        "--output",
-        default=None,
-        help="figure path, <input>/relaxation_time.png if unset",
-    )
-    parser.add_argument("--confidence", type=float, default=0.95)
-    parser.add_argument(
-        "--extrapolate-sequential",
-        type=int,
-        default=None,
-        metavar="BATCH_SIZE",
-        help="extend the sequential curve as a dashed linear fit up to this batch size",
-    )
-    parser.add_argument(
-        "--paper",
-        action="store_true",
-        help="single-column figure without title, for the manuscript",
-    )
-    parser.add_argument("--show", action="store_true", help="also open the figure")
-    return parser.parse_args()
-
-
-def load_timings(input_dir):
-    """The run settings and, per method, its batch sizes and the wall times as a
-    ``(n_batch_sizes, n_repeats)`` array.
-
-    A method only covers the batch sizes it was timed at: the sequential baseline may
-    stop early (``--sequential-max-batch-size``), and a run that died keeps the batch
-    sizes it finished.
-    """
-    with open(os.path.join(input_dir, "timings.json")) as f:
-        results = json.load(f)
-    meta = results["meta"]
-
-    timings = {}
-    for method in METHODS:
-        records = [r for r in results["records"] if r["method"] == method]
-        batch_sizes = sorted({r["batch_size"] for r in records})
-        samples = np.full((len(batch_sizes), meta["repeats"]), np.nan)
-        for record in records:
-            size_idx = batch_sizes.index(record["batch_size"])
-            samples[size_idx, record["repeat"]] = record["time_s"]
-        assert not np.isnan(samples).any(), f"{method}: missing timings"
-        timings[method] = (batch_sizes, samples)
-    return meta, timings
-
-
-def mean_ci(samples, confidence):
-    """Mean and half-width of the Student-t confidence interval, along the repeats."""
-    mean = samples.mean(axis=1)
-    half_width = stats.sem(samples, axis=1) * stats.t.ppf(
-        0.5 + confidence / 2, samples.shape[1] - 1
-    )
-    return mean, half_width
-
-
-def linear_fit(batch_sizes, samples):
-    """Intercept and slope of ``time = a + b * batch_size``, fitted to every repeat."""
-    sizes = np.repeat(batch_sizes, samples.shape[1])
-    slope, intercept = np.polyfit(sizes, samples.ravel(), 1)
-    return intercept, slope
-
-
-def extrapolate(batch_sizes, samples, all_batch_sizes, max_batch_size):
-    """The linear fit evaluated from the last timed batch size up to ``max_batch_size``,
-    at the batch sizes of the run and at ``max_batch_size`` itself."""
-    intercept, slope = linear_fit(batch_sizes, samples)
-    sizes = [n for n in all_batch_sizes if batch_sizes[-1] <= n < max_batch_size]
-    sizes = np.array([batch_sizes[-1], *sizes[1:], max_batch_size])
-    return sizes, intercept + slope * sizes
-
-
-def print_table(summary):
-    # (mean, half-width) per method and batch size, absent where it was not timed
-    rows = {
-        method: {n: (m, ci) for n, m, ci in zip(sizes, mean, half_width)}
-        for method, (sizes, mean, half_width) in summary.items()
-    }
-    header = f"{'batch size':>10} {'batch-wise [s]':>22} {'sequential [s]':>22} {'speedup':>9}"
-    print(header)
-    print("-" * len(header))
-    for n_structures in sorted(set().union(*rows.values())):
-        line = f"{n_structures:>10}"
-        for method in METHODS:
-            if n_structures in rows[method]:
-                mean, half_width = rows[method][n_structures]
-                line += f" {mean:>13.2f} +- {half_width:<6.2f}"
-            else:
-                line += f" {'-':>22}"
-        if all(n_structures in rows[method] for method in METHODS):
-            speedup = (
-                rows["sequential"][n_structures][0]
-                / rows["batch-wise"][n_structures][0]
-            )
-            line += f" {speedup:>8.2f}x"
-        print(line)
-
-
-def plot(meta, summary, confidence, extrapolation=None, layout="screen"):
-    """``extrapolation`` is an optional ``(sizes, times)`` continuation of the
-    sequential curve, drawn dashed. ``layout`` is a key of ``LAYOUTS``."""
-    style = LAYOUTS[layout]
-    batch_sizes = sorted(set().union(*(sizes for sizes, *_ in summary.values())))
-    if extrapolation is not None:
-        batch_sizes = sorted({*batch_sizes, *extrapolation[0].tolist()})
-
-    fig, ax = plt.subplots(figsize=style["figsize"], dpi=style["dpi"])
-    fig.patch.set_facecolor(style["surface"])
-    ax.set_facecolor(style["surface"])
-
-    for label, (sizes, mean, half_width) in summary.items():
-        color = SERIES[label]
-        # confidence interval as a band, the mean as the line on top of it
-        ax.fill_between(
-            sizes,
-            mean - half_width,
-            mean + half_width,
-            color=color,
-            alpha=0.18,
-            linewidth=0,
-            zorder=2,
-        )
-        ax.plot(
-            sizes,
-            mean,
-            color=color,
-            linewidth=style["line"],
-            marker="o",
-            markersize=style["marker"],
-            markeredgecolor=style["surface"],
-            markeredgewidth=style["marker"] / 3,
-            label=label,
-            zorder=3,
-        )
-        end = (sizes[-1], mean[-1])
-        if label == "sequential" and extrapolation is not None:
-            ax.plot(
-                *extrapolation,
-                color=color,
-                linewidth=0.75 * style["line"],
-                linestyle=(0, (4, 3)),
-                label="linear fit",
-                zorder=3,
-            )
-            end = (extrapolation[0][-1], extrapolation[1][-1])
-        # label the series at its end, so the curves can be read without the legend;
-        # a single column has no room for it to the right, there the legend does it
-        if style["title"]:
-            ax.annotate(
-                label,
-                end,
-                textcoords="offset points",
-                xytext=(10, 0),
-                va="center",
-                color=color,
-                fontsize=style["font"],
-            )
-
-    if extrapolation is not None:
-        # the gap between both curves at the end of the extrapolation, as a speedup
-        n_end, sequential_end = extrapolation[0][-1], extrapolation[1][-1]
-        sizes, mean, _ = summary["batch-wise"]
-        if n_end in sizes:
-            batchwise_end = mean[sizes.index(n_end)]
-            ax.annotate(
-                "",
-                (n_end, batchwise_end),
-                xytext=(n_end, sequential_end),
-                arrowprops=dict(
-                    arrowstyle="<->",
-                    color=INK_SECONDARY,
-                    linewidth=style["line"] / 2,
-                    shrinkA=style["marker"],
-                    shrinkB=style["marker"],
-                ),
-                zorder=4,
-            )
-            ax.annotate(
-                f"~{sequential_end / batchwise_end:.0f}\u00d7"
-                + ("\n" if layout == "paper" else " ")
-                + "faster",
-                (n_end, np.sqrt(sequential_end * batchwise_end)),  # middle on log axes
-                # left of the arrow, where neither curve passes
-                textcoords="offset points",
-                xytext=(-style["pad"], 0),
-                ha="right",
-                va="center",
-                color=INK,
-                fontsize=style["font"],
-                fontweight="bold",
-            )
-
-    # both scalings are power laws, which read as straight lines on log-log axes
-    ax.set_xscale("log", base=2)
-    ax.set_yscale("log", base=2)
-    ax.set_xticks(batch_sizes)
-    ax.set_xticklabels(
-        [n if i % style["tick_every"] == 0 else "" for i, n in enumerate(batch_sizes)]
-    )
-    ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda y, _: f"{y:g}"))
-    ax.minorticks_off()
-    ax.set_xlim(
-        batch_sizes[0] / 1.3, batch_sizes[-1] * (2.6 if style["title"] else 1.3)
-    )
-
-    ax.grid(True, which="major", color=GRID, linewidth=0.4 * style["line"])
-    ax.set_axisbelow(True)
-    for side in ("top", "right"):
-        ax.spines[side].set_visible(False)
-    for side in ("left", "bottom"):
-        ax.spines[side].set_color(INK)
-    ax.tick_params(colors=INK, labelsize=style["tick_font"], length=0)
-
-    ax.set_xlabel(
-        "batch size (number of structures)",
-        color=INK,
-        fontsize=style["font"],
-        labelpad=style["pad"],
-    )
-    ax.set_ylabel(
-        "relaxation time [s]",
-        color=INK,
-        fontsize=style["font"],
-        labelpad=style["pad"],
-    )
-    legend = ax.legend(frameon=False, loc="upper left", fontsize=style["font"])
-    for text in legend.get_texts():
-        text.set_color(INK)
-
-    if not style["title"]:
-        fig.tight_layout(pad=0.3)
-        return fig
-
-    ax.set_title(
-        "Batch-wise relaxation pulls ahead as the batch grows",
-        color=INK,
-        fontsize=13,
-        loc="left",
-        pad=34,
-    )
-    ax.text(
-        0.0,
-        1.015,
-        f"wall time to relax the whole batch to fmax = {meta['fmax']} eV/A "
-        f"on {meta.get('gpu') or meta['device']},\n"
-        f"mean of {meta['repeats']} runs, band is the {confidence:.0%} "
-        "confidence interval",
-        transform=ax.transAxes,
-        color=MUTED,
-        fontsize=9,
-        va="bottom",
-        linespacing=1.6,
-    )
-    fig.tight_layout()
-    return fig
-
-
-def main():
-    args = parse_args()
-    meta, timings = load_timings(args.input)
-    summary = {
-        method: (sizes, *mean_ci(samples, args.confidence))
-        for method, (sizes, samples) in timings.items()
-    }
-
-    print_table(summary)
-
-    extrapolation = None
-    if args.extrapolate_sequential is not None:
-        sizes, samples = timings["sequential"]
-        intercept, slope = linear_fit(sizes, samples)
-        print(f"sequential linear fit: {intercept:.3f} s + {slope:.4f} s * batch size")
-        extrapolation = extrapolate(
-            sizes, samples, meta["batch_sizes"], args.extrapolate_sequential
-        )
-
-    fig = plot(
-        meta,
-        summary,
-        args.confidence,
-        extrapolation,
-        layout="paper" if args.paper else "screen",
-    )
-    output = args.output or os.path.join(args.input, "relaxation_time.png")
-    fig.savefig(output, facecolor=fig.get_facecolor())
-    print(f"figure written to {output}")
-    if args.show:
-        plt.show()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/benchmarks/batchwise_relaxation/run_benchmark.py b/examples/benchmarks/batchwise_relaxation/run_benchmark.py
deleted file mode 100644
index 3c7ef1e5f..000000000
--- a/examples/benchmarks/batchwise_relaxation/run_benchmark.py
+++ /dev/null
@@ -1,258 +0,0 @@
-"""Time the batch-wise ``dynamics.LBFGS`` against sequential ase ``LBFGS``.
-
-Both optimizers relax the same batches of noisy structures over a range of batch sizes.
-The relaxed structures and the wall times are written to the output directory, which
-``plot_benchmark.py`` reads to plot the relaxation time against the batch size.
-
-Only the relaxation itself is timed: building the calculator, converting the structures
-and constructing the relaxer all happen before the clock starts. Every batch size is
-repeated with a different noise seed per repeat, and both optimizers relax the identical
-batch in each repeat (paired samples).
-"""
-
-import argparse
-import json
-import os
-import socket
-import time
-from copy import deepcopy
-
-import numpy as np
-import torch
-from ase.io import read, write
-from ase.optimize import LBFGS
-
-import schnetpack as spk
-from schnetpack.dynamics import ForceCalculator
-from schnetpack.interfaces.ase_interface import (
-    SpkCalculator,
-    atoms_to_batch,
-    batch_to_atoms,
-)
-from schnetpack.utils.compatibility import load_model
-
-TESTDATA = os.path.join(
-    os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "tests", "testdata"
-)
-METHODS = ("batch-wise", "sequential")
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
-    parser.add_argument(
-        "--model",
-        default=os.path.join(TESTDATA, "md_ethanol.model"),
-        help="SchNetPack model that predicts energies and forces",
-    )
-    parser.add_argument(
-        "--structures",
-        default=os.path.join(TESTDATA, "ethanol_conformers.xyz"),
-        help="starting structures, cycled through to fill a batch",
-    )
-    parser.add_argument(
-        "--batch-sizes", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32]
-    )
-    parser.add_argument("--repeats", type=int, default=5)
-    parser.add_argument(
-        "--sequential-max-batch-size",
-        type=int,
-        default=None,
-        help="time the sequential baseline only up to this batch size, it scales "
-        "linearly and dominates the run time of large batches; no limit if unset",
-    )
-    parser.add_argument(
-        "--fmax", type=float, default=0.001, help="convergence criterion in eV/Ang"
-    )
-    parser.add_argument("--max-steps", type=int, default=1000)
-    parser.add_argument(
-        "--noise",
-        type=float,
-        default=0.05,
-        help="standard deviation of the Gaussian noise on the positions, in Ang",
-    )
-    parser.add_argument("--cutoff-skin", type=float, default=0.3)
-    parser.add_argument("--energy-unit", default="kcal/mol")
-    parser.add_argument("--position-unit", default="Ang")
-    parser.add_argument(
-        "--device", default="cuda" if torch.cuda.is_available() else "cpu"
-    )
-    parser.add_argument("--output", default="./batchwise_relaxation_benchmark")
-    return parser.parse_args()
-
-
-def make_structures(conformers, n_structures, noise, seed=None):
-    """``n_structures`` starting geometries, cycling through the conformers.
-
-    Every copy gets Gaussian noise on its atomic positions, so the structures of a batch
-    are similar but not identical. Pass a ``seed`` to get a reproducible set.
-    """
-    rng = np.random.default_rng(seed)
-    structures = []
-    for idx in range(n_structures):
-        noisy_atoms = conformers[idx % len(conformers)].copy()
-        noisy_atoms.positions += rng.normal(
-            scale=noise, size=noisy_atoms.positions.shape
-        )
-        structures.append(noisy_atoms)
-    return structures
-
-
-class Benchmark:
-    """Builds both optimizers from the same settings and times their relaxations."""
-
-    def __init__(self, args):
-        self.args = args
-        self.device = torch.device(args.device)
-        model = load_model(args.model, device=self.device)
-        self.cutoff = model.representation.cutoff.item()
-
-    def batchwise_relaxer(self):
-        """An ``LBFGS`` that relaxes a whole batch in a single model call per step."""
-        calculator = ForceCalculator(
-            self.args.model,
-            neighbor_list=spk.dynamics.BatchNeighborList(
-                neighbor_list=spk.transform.MatScipyNeighborList(cutoff=self.cutoff),
-                cutoff_skin=self.args.cutoff_skin,
-            ),
-            device=self.device,
-            dtype=torch.float32,
-            energy_unit=self.args.energy_unit,
-            position_unit=self.args.position_unit,
-        )
-        return spk.dynamics.LBFGS(calculator, fmax=self.args.fmax)
-
-    def ase_calculator(self):
-        """A plain ase calculator, which evaluates one structure per call."""
-        return SpkCalculator(
-            model=self.args.model,
-            neighbor_list=spk.transform.MatScipyNeighborList(cutoff=self.cutoff),
-            device=self.device,
-            energy_unit=self.args.energy_unit,
-            position_unit=self.args.position_unit,
-        )
-
-    def time_batchwise(self, structures):
-        """Wall time of relaxing the whole batch in parallel. The setup is not timed."""
-        relaxer = self.batchwise_relaxer()
-        inputs = relaxer.calculator.prepare(
-            atoms_to_batch(structures, device=self.device)
-        )
-
-        start = time.perf_counter()
-        relaxed = relaxer.run(inputs, self.args.max_steps)
-        elapsed = time.perf_counter() - start
-
-        # a run that is fast because it never converged is not a faster run
-        max_force = relaxer.calculator(relaxed)["forces"].norm(dim=-1).max()
-        assert max_force < self.args.fmax, "batch-wise run hit the step limit"
-        return elapsed, batch_to_atoms(relaxed)
-
-    def time_sequential(self, structures):
-        """Wall time of relaxing the same structures one after another."""
-        calculator = self.ase_calculator()
-        structures = deepcopy(structures)  # LBFGS relaxes the Atoms objects in place
-
-        start = time.perf_counter()
-        steps = []
-        for structure in structures:
-            structure.calc = calculator
-            lbfgs = LBFGS(structure, logfile=None)
-            lbfgs.run(fmax=self.args.fmax, steps=self.args.max_steps)
-            steps.append(lbfgs.nsteps)
-        elapsed = time.perf_counter() - start
-
-        assert max(steps) < self.args.max_steps, "a sequential run hit the step limit"
-        # the relaxed positions only, without the calculator attached
-        for structure in structures:
-            structure.calc = None
-        return elapsed, structures
-
-
-def main():
-    args = parse_args()
-    structures_dir = os.path.join(args.output, "structures")
-    os.makedirs(structures_dir, exist_ok=True)
-
-    conformers = read(args.structures, index=":")
-    benchmark = Benchmark(args)
-    timers = {
-        "batch-wise": benchmark.time_batchwise,
-        "sequential": benchmark.time_sequential,
-    }
-
-    # The first relaxation pays for lazy torch initialisation, which would otherwise
-    # land entirely on the smallest batch size and fake a slow start.
-    warmup = make_structures(conformers, 2, args.noise, seed=999)
-    for timer in timers.values():
-        timer(warmup)
-
-    meta = dict(
-        model=os.path.abspath(args.model),
-        structures=os.path.abspath(args.structures),
-        device=benchmark.device.type,
-        # the hardware, so runs on different nodes can be told apart
-        gpu=(
-            torch.cuda.get_device_name(benchmark.device)
-            if benchmark.device.type == "cuda"
-            else None
-        ),
-        hostname=socket.gethostname(),
-        slurm_job_id=os.environ.get("SLURM_JOB_ID"),
-        slurm_partition=os.environ.get("SLURM_JOB_PARTITION"),
-        batch_sizes=args.batch_sizes,
-        repeats=args.repeats,
-        sequential_max_batch_size=args.sequential_max_batch_size,
-        fmax=args.fmax,
-        max_steps=args.max_steps,
-        noise=args.noise,
-        cutoff_skin=args.cutoff_skin,
-        torch_version=torch.__version__,
-        schnetpack_version=spk.__version__,
-    )
-    timings_file = os.path.join(args.output, "timings.json")
-
-    records = []
-    for n_structures in args.batch_sizes:
-        methods = [
-            method
-            for method in METHODS
-            if method != "sequential"
-            or args.sequential_max_batch_size is None
-            or n_structures <= args.sequential_max_batch_size
-        ]
-        times = {method: [] for method in methods}
-        for repeat in range(args.repeats):
-            # both optimizers relax the identical batch, so the repeats are paired
-            structures = make_structures(conformers, n_structures, args.noise, repeat)
-            for method in methods:
-                elapsed, relaxed = timers[method](structures)
-                structures_file = os.path.join(
-                    "structures", f"{method}_n{n_structures:03d}_r{repeat:02d}.xyz"
-                )
-                write(os.path.join(args.output, structures_file), relaxed)
-                records.append(
-                    dict(
-                        method=method,
-                        batch_size=n_structures,
-                        repeat=repeat,
-                        seed=repeat,
-                        time_s=elapsed,
-                        structures_file=structures_file,
-                    )
-                )
-                times[method].append(elapsed)
-
-        print(
-            f"batch size {n_structures:>5}: "
-            + ", ".join(f"{m} {np.mean(t):6.2f} s" for m, t in times.items()),
-            flush=True,
-        )
-        # after every batch size, so a run that dies at a large batch keeps the rest
-        with open(timings_file, "w") as f:
-            json.dump(dict(meta=meta, records=records), f, indent=2)
-
-    print(f"results written to {timings_file}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/benchmarks/batchwise_relaxation/submit_benchmark.sh b/examples/benchmarks/batchwise_relaxation/submit_benchmark.sh
deleted file mode 100755
index 1bce95601..000000000
--- a/examples/benchmarks/batchwise_relaxation/submit_benchmark.sh
+++ /dev/null
@@ -1,33 +0,0 @@
-#!/bin/bash -l
-#SBATCH --partition=gpu-test
-#SBATCH --gres=gpu:80gb:1
-# Check the H100 GRES name with: sinfo -o "%P %G %f" (gpu:80gb may also match an A100)
-# Whole node, so other jobs do not distort the wall times
-#SBATCH --exclusive
-# Standard output and error:
-#SBATCH -o job_%j.out        # Standard output, %j = job ID
-#SBATCH -e job_%j.err        # Standard error, %j = job ID
-# Initial working directory (submit from this folder):
-#SBATCH -D ./
-# Job Name:
-#SBATCH -J relax_benchmark
-# Number of nodes
-#SBATCH --nodes=1
-#SBATCH --cpus-per-task=4
-
-# Time batch-wise dynamics.LBFGS against sequential ase LBFGS.
-# Submit from this folder: sbatch submit_benchmark.sh
-# Plot afterwards with plot_benchmark.py.
-
-CONTAINER=/home/Jonas/container/schnetpack/schnetpack_container.sif
-OUTPUT=$SLURM_SUBMIT_DIR/results_$SLURM_JOB_ID
-
-# pin the torch CPU threads to the allocated cores
-export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
-
-apptainer exec --nv $CONTAINER python $SLURM_SUBMIT_DIR/run_benchmark.py \
-  --device cuda \
-  --batch-sizes 1 2 4 8 16 32 64 128 256 512 1024 2048 4096 8192 \
-  --repeats 5 \
-  --sequential-max-batch-size 32 \
-  --output $OUTPUT
diff --git a/examples/benchmarks/distillation/plot_inference_benchmark.py b/examples/benchmarks/distillation/plot_inference_benchmark.py
deleted file mode 100644
index 2f3cf9411..000000000
--- a/examples/benchmarks/distillation/plot_inference_benchmark.py
+++ /dev/null
@@ -1,187 +0,0 @@
-"""Plot the inference times written by ``run_inference_benchmark.py`` over batch size.
-
-The figure shows the wall time of one energy and force call, as the median over the
-repeats with the interquartile range as error bars. A model that ran out of GPU memory
-has no point at that batch size.
-"""
-
-import argparse
-import json
-import os
-
-import matplotlib
-
-matplotlib.use("Agg")
-import matplotlib.pyplot as plt
-import matplotlib.ticker as ticker
-import numpy as np
-
-# categorical slots in fixed order, so a model keeps its color across figures
-COLORS = ("#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300")
-MARKERS = ("o", "s", "D", "^", "v", "P")
-INK = "#000000"
-GRID = "#e4e3df"
-DEFAULT_TIMINGS = (
-    "/home/jonas/Documents/cluster/schnetpack_debugging/distillation/"
-    "results_5124647/timings.json"
-)
-# the names of the models in the manuscript, as in its table of the distillation
-PAPER_NAMES = {
-    "teacher": "teacher (MACE-OFF23)",
-    "student_f": "student, E/F",
-    "student_hvp": "student, E/F/HVP",
-}
-
-# "screen" to inspect a run, "paper" for a single column of a two-column journal: no
-# title (the caption carries it), print font sizes
-LAYOUTS = {
-    "screen": dict(
-        figsize=(6, 4.5),
-        font=10,
-        tick_font=10,
-        line=2,
-        marker=7,
-        title=True,
-        extensions=("png", "pdf"),
-    ),
-    "paper": dict(
-        figsize=(3.37, 2.5),
-        font=8,
-        tick_font=7,
-        line=1.3,
-        marker=4,
-        title=False,
-        extensions=("pdf",),
-    ),
-}
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
-    parser.add_argument(
-        "timings",
-        nargs="?",
-        default=DEFAULT_TIMINGS,
-        help="the timings.json of a benchmark run",
-    )
-    parser.add_argument(
-        "--output",
-        default=None,
-        help="the figure, without extension; a png and a pdf are written, only a pdf "
-        "with --paper. Defaults to inference_vs_batch_size next to the timings",
-    )
-    parser.add_argument(
-        "--paper",
-        action="store_true",
-        help="single-column figure without title, for the manuscript",
-    )
-    return parser.parse_args()
-
-
-def summarize(records):
-    """``{model: {batch_size: record}}`` with the median and quartiles of the times."""
-    models = {}
-    for record in records:
-        record = dict(record)
-        if not record["oom"]:
-            times = np.asarray(record["times_s"])
-            record["median_s"] = np.median(times)
-            record["q25_s"], record["q75_s"] = np.percentile(times, [25, 75])
-        models.setdefault(record["model"], {})[record["batch_size"]] = record
-    return models
-
-
-def style_axis(ax, batch_sizes, style):
-    ax.set_xscale("log")
-    ax.set_yscale("log")
-    ax.set_xticks(batch_sizes)
-    ax.set_xticklabels([str(n) for n in batch_sizes])
-    ax.xaxis.set_minor_locator(ticker.NullLocator())
-    # decades alone leave a panel spanning less than one with a single label
-    ax.yaxis.set_major_locator(ticker.LogLocator(subs=(1, 2, 5)))
-    ax.yaxis.set_major_formatter(ticker.FormatStrFormatter("%g"))
-    ax.yaxis.set_minor_locator(ticker.NullLocator())
-    ax.set_xlabel("batch size [structures]", color=INK, fontsize=style["font"])
-    ax.grid(True, which="major", color=GRID, linewidth=0.4 * style["line"])
-    ax.set_axisbelow(True)
-    for side in ("top", "right"):
-        ax.spines[side].set_visible(False)
-    for side in ("left", "bottom"):
-        ax.spines[side].set_color(INK)
-    ax.tick_params(colors=INK, labelsize=style["tick_font"])
-
-
-def plot(models, meta, output, layout="screen"):
-    style = LAYOUTS[layout]
-    batch_sizes = meta["batch_sizes"]
-    fig, ax = plt.subplots(figsize=style["figsize"])
-
-    for idx, (label, by_size) in enumerate(models.items()):
-        color, marker = COLORS[idx % len(COLORS)], MARKERS[idx % len(MARKERS)]
-        done = [n for n in batch_sizes if n in by_size and not by_size[n]["oom"]]
-        median = np.array([by_size[n]["median_s"] for n in done])
-        q25 = np.array([by_size[n]["q25_s"] for n in done])
-        q75 = np.array([by_size[n]["q75_s"] for n in done])
-        if layout == "paper":
-            name = PAPER_NAMES.get(label, label)
-        else:
-            n_parameters = meta["models"][label]["n_parameters"]
-            name = f"{label} ({n_parameters / 1e6:.2f}M params)"
-
-        ax.errorbar(
-            done,
-            1e3 * median,
-            yerr=1e3 * np.stack([median - q25, q75 - median]),
-            color=color,
-            marker=marker,
-            markersize=style["marker"],
-            linewidth=style["line"],
-            capsize=style["marker"] / 2,
-            label=name,
-        )
-
-    style_axis(ax, batch_sizes, style)
-    ax.set_ylabel(
-        "time per energy + force call [ms]", color=INK, fontsize=style["font"]
-    )
-    ax.legend(frameon=False, fontsize=style["tick_font"], loc="upper left")
-    if style["title"]:
-        ax.set_title(
-            f"Inference on {meta['gpu'] or meta['device']}, {meta['dtype']}, "
-            f"{meta['repeats']} repeats (job {meta['slurm_job_id']})",
-            color=INK,
-            fontsize=style["font"],
-        )
-    fig.tight_layout(pad=0.3 if layout == "paper" else 1.08)
-    paths = [f"{output}.{extension}" for extension in style["extensions"]]
-    for path in paths:
-        fig.savefig(path, dpi=300 if layout == "paper" else 200)
-    plt.close(fig)
-    return paths
-
-
-def main():
-    args = parse_args()
-    with open(args.timings) as f:
-        timings = json.load(f)
-    output = args.output or os.path.join(
-        os.path.dirname(os.path.abspath(args.timings)), "inference_vs_batch_size"
-    )
-    output = os.path.splitext(output)[0] if output.endswith(".pdf") else output
-
-    models = summarize(timings["records"])
-    for label, by_size in models.items():
-        line = ", ".join(
-            f"{n}: OOM" if record["oom"] else f"{n}: {1e3 * record['median_s']:.2f} ms"
-            for n, record in sorted(by_size.items())
-        )
-        print(f"{label}: {line}")
-
-    paths = plot(
-        models, timings["meta"], output, layout="paper" if args.paper else "screen"
-    )
-    print(f"figure written to {' and '.join(paths)}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/benchmarks/distillation/run_inference_benchmark.py b/examples/benchmarks/distillation/run_inference_benchmark.py
deleted file mode 100644
index 496ef49fc..000000000
--- a/examples/benchmarks/distillation/run_inference_benchmark.py
+++ /dev/null
@@ -1,317 +0,0 @@
-"""Time the inference of a distillation teacher against its students.
-
-Every model predicts energies and forces for the same batches of structures over a range
-of batch sizes. The wall times are written to ``timings.json`` in the output directory.
-
-Only the model call is timed. The batch, with its neighbor list at the model's own
-cutoff, is built before the clock starts; how long that build takes is recorded
-separately, since a student with a smaller cutoff saves there as well. Within each batch
-size the models take turns, so that drifts of the node (clocks, temperature) hit all of
-them alike. A model that runs out of GPU memory is recorded as such and skips the larger
-batch sizes.
-"""
-
-import argparse
-import json
-import os
-import socket
-import time
-
-import numpy as np
-import torch
-from ase.io import read
-
-import schnetpack as spk
-from schnetpack.dynamics import ForceCalculator
-from schnetpack.interfaces.ase_interface import AtomsConverter
-from schnetpack.model import MaceModel
-from schnetpack.model.mace import load_mace
-from schnetpack.utils.compatibility import load_model
-
-TESTDATA = os.path.join(
-    os.path.dirname(os.path.abspath(__file__)), "..", "..", "tests", "testdata"
-)
-MODEL_FORMATS = ("mace", "schnetpack")
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
-    parser.add_argument(
-        "--teacher",
-        required=True,
-        help="the teacher: a MACE TorchScript archive or a SchNetPack model",
-    )
-    parser.add_argument(
-        "--teacher-format",
-        choices=MODEL_FORMATS,
-        default="mace",
-        help="how the teacher is loaded",
-    )
-    parser.add_argument(
-        "--students",
-        nargs="+",
-        default=[],
-        metavar="[LABEL=]PATH",
-        help="SchNetPack models distilled from the teacher, e.g. the best_model of a "
-        "spktrain run; an optional label names them in the results",
-    )
-    parser.add_argument(
-        "--structures",
-        default=os.path.join(TESTDATA, "ethanol_conformers.xyz"),
-        help="structures in any format ase reads, cycled through to fill a batch",
-    )
-    parser.add_argument(
-        "--batch-sizes", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32, 64, 128]
-    )
-    parser.add_argument(
-        "--repeats", type=int, default=20, help="timed model calls per batch size"
-    )
-    parser.add_argument(
-        "--warmup", type=int, default=5, help="untimed model calls per batch size"
-    )
-    parser.add_argument(
-        "--noise",
-        type=float,
-        default=0.05,
-        help="standard deviation of the Gaussian noise on the positions, in Ang",
-    )
-    parser.add_argument("--dtype", choices=("float32", "float64"), default="float32")
-    parser.add_argument(
-        "--device", default="cuda" if torch.cuda.is_available() else "cpu"
-    )
-    parser.add_argument("--output", default="./inference_benchmark")
-    return parser.parse_args()
-
-
-def make_structures(conformers, n_structures, noise, seed=None):
-    """``n_structures`` geometries, cycling through the conformers.
-
-    Every copy gets Gaussian noise on its atomic positions, so the structures of a batch
-    are similar but not identical. Pass a ``seed`` to get a reproducible set.
-    """
-    rng = np.random.default_rng(seed)
-    structures = []
-    for idx in range(n_structures):
-        noisy_atoms = conformers[idx % len(conformers)].copy()
-        noisy_atoms.positions += rng.normal(
-            scale=noise, size=noisy_atoms.positions.shape
-        )
-        structures.append(noisy_atoms)
-    return structures
-
-
-def parse_students(entries):
-    """``{label: path}`` from ``[LABEL=]PATH`` entries."""
-    students = {}
-    for idx, entry in enumerate(entries):
-        label, sep, path = entry.partition("=")
-        if not sep:
-            label, path = f"student_{idx}", entry
-        if label in students or label == "teacher":
-            raise ValueError(f"the label {label!r} is used twice")
-        students[label] = path
-    return students
-
-
-class TimedModel:
-    """A model loaded for inference, with the cutoff its neighbor list needs."""
-
-    def __init__(self, label, path, model_format, device, dtype):
-        self.label = label
-        self.path = os.path.abspath(path)
-        self.model_format = model_format
-        if model_format == "mace":
-            # the core alone is cast, so the E0s stay in float64
-            model = MaceModel(load_mace(path), dtype=dtype)
-            self.cutoff = model.r_max
-        else:
-            model = load_model(path, device="cpu").to(dtype)
-            self.cutoff = model.representation.cutoff.item()
-        model.eval()
-        for parameter in model.parameters():
-            parameter.requires_grad_(False)
-        self.n_parameters = sum(p.numel() for p in model.parameters())
-        self.calculator = ForceCalculator(model, device=device)
-        self.converter = AtomsConverter(
-            neighbor_list=spk.transform.MatScipyNeighborList(cutoff=self.cutoff),
-            device=device,
-            dtype=dtype,
-        )
-
-    def meta(self):
-        return dict(
-            path=self.path,
-            format=self.model_format,
-            cutoff=self.cutoff,
-            n_parameters=self.n_parameters,
-        )
-
-
-def is_out_of_memory(err):
-    """Whether ``err`` is CUDA running out of memory. Inside a TorchScript model, such
-    as a MACE archive, it is raised as a plain ``RuntimeError``."""
-    return isinstance(err, torch.OutOfMemoryError) or "CUDA out of memory" in str(err)
-
-
-def synchronize(device):
-    if device.type == "cuda":
-        torch.cuda.synchronize(device)
-
-
-def memory_mb(device, *, peak=False, reserved=False):
-    """The GPU memory torch has allocated, or reserved in its caching allocator, now or
-    at its peak, in MiB; None on the cpu."""
-    if device.type != "cuda":
-        return None
-    if reserved:
-        memory = torch.cuda.max_memory_reserved if peak else torch.cuda.memory_reserved
-    else:
-        memory = (
-            torch.cuda.max_memory_allocated if peak else torch.cuda.memory_allocated
-        )
-    return memory(device) / 2**20
-
-
-def time_model(model, structures, args, device):
-    """Wall times of ``args.repeats`` energy and force calls on one batch, after
-    ``args.warmup`` untimed ones, together with the time the batch took to build and
-    the GPU memory of the calls.
-
-    Every model stays on the device for the whole run, so all memory figures include
-    the weights of all of them: ``resident_memory_mb`` is what is allocated before the
-    batch is built, ``allocated_memory_mb`` what is allocated with the batch in place,
-    before the calls. ``peak_memory_mb`` is the most allocated during the timed calls,
-    and ``peak_reserved_memory_mb`` the most the caching allocator held then, which is
-    closer to what ``nvidia-smi`` shows.
-    """
-    if device.type == "cuda":
-        # the cache still holds the blocks of the previous model's batch
-        torch.cuda.empty_cache()
-    resident_memory_mb = memory_mb(device)
-    start = time.perf_counter()
-    batch = model.converter(structures)
-    synchronize(device)
-    neighbor_list_s = time.perf_counter() - start
-    allocated_memory_mb = memory_mb(device)
-
-    for _ in range(args.warmup):
-        model.calculator(batch)
-    synchronize(device)
-
-    if device.type == "cuda":
-        torch.cuda.reset_peak_memory_stats(device)
-    times = []
-    for _ in range(args.repeats):
-        start = time.perf_counter()
-        outputs = model.calculator(batch)
-        synchronize(device)
-        times.append(time.perf_counter() - start)
-
-    # a call that is fast because it returned nonsense is not a faster call
-    assert torch.isfinite(outputs["forces"]).all(), f"{model.label}: non-finite forces"
-    return dict(
-        times_s=times,
-        neighbor_list_s=neighbor_list_s,
-        resident_memory_mb=resident_memory_mb,
-        allocated_memory_mb=allocated_memory_mb,
-        peak_memory_mb=memory_mb(device, peak=True),
-        peak_reserved_memory_mb=memory_mb(device, peak=True, reserved=True),
-    )
-
-
-def main():
-    args = parse_args()
-    os.makedirs(args.output, exist_ok=True)
-    device = torch.device(args.device)
-    dtype = getattr(torch, args.dtype)
-
-    models = [
-        TimedModel("teacher", args.teacher, args.teacher_format, device, dtype)
-    ] + [
-        TimedModel(label, path, "schnetpack", device, dtype)
-        for label, path in parse_students(args.students).items()
-    ]
-    conformers = read(args.structures, index=":")
-
-    # The first model call pays for lazy CUDA initialisation (cuBLAS workspace and
-    # the like), which stays allocated: it would land on the first model measured,
-    # in both its time and its memory.
-    warmup = make_structures(conformers, 2, args.noise, seed=999)
-    for model in models:
-        model.calculator(model.converter(warmup))
-    synchronize(device)
-
-    meta = dict(
-        models={model.label: model.meta() for model in models},
-        structures=os.path.abspath(args.structures),
-        device=device.type,
-        # the hardware, so runs on different nodes can be told apart
-        gpu=torch.cuda.get_device_name(device) if device.type == "cuda" else None,
-        hostname=socket.gethostname(),
-        slurm_job_id=os.environ.get("SLURM_JOB_ID"),
-        slurm_partition=os.environ.get("SLURM_JOB_PARTITION"),
-        cpu_threads=torch.get_num_threads(),
-        batch_sizes=args.batch_sizes,
-        repeats=args.repeats,
-        warmup=args.warmup,
-        noise=args.noise,
-        dtype=args.dtype,
-        torch_version=torch.__version__,
-        schnetpack_version=spk.__version__,
-    )
-    timings_file = os.path.join(args.output, "timings.json")
-
-    records = []
-    out_of_memory = set()
-    for n_structures in args.batch_sizes:
-        structures = make_structures(conformers, n_structures, args.noise, seed=0)
-        n_atoms = sum(len(structure) for structure in structures)
-        summary = []
-        for model in models:
-            if model.label in out_of_memory:
-                continue
-            try:
-                result = time_model(model, structures, args, device)
-            except RuntimeError as err:
-                if not is_out_of_memory(err):
-                    raise
-                result = None
-            if result is None:
-                # outside the except block, whose traceback holds the failed batch
-                torch.cuda.empty_cache()
-                out_of_memory.add(model.label)
-                records.append(
-                    dict(model=model.label, batch_size=n_structures, oom=True)
-                )
-                summary.append(f"{model.label} out of memory")
-                continue
-            records.append(
-                dict(
-                    model=model.label,
-                    batch_size=n_structures,
-                    n_atoms=n_atoms,
-                    oom=False,
-                    **result,
-                )
-            )
-            median = np.median(result["times_s"])
-            line = (
-                f"{model.label} {1e3 * median:8.2f} ms "
-                f"({n_structures / median:9.1f} structures/s"
-            )
-            if result["peak_memory_mb"] is not None:
-                # the model's own share: the weights of all models are resident
-                own_memory_mb = result["peak_memory_mb"] - result["resident_memory_mb"]
-                line += f", {own_memory_mb:8.1f} MiB"
-            summary.append(line + ")")
-
-        print(f"batch size {n_structures:>5}: " + ", ".join(summary), flush=True)
-        # after every batch size, so a run that dies at a large batch keeps the rest
-        with open(timings_file, "w") as f:
-            json.dump(dict(meta=meta, records=records), f, indent=2)
-
-    print(f"results written to {timings_file}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/benchmarks/distillation/submit_inference_benchmark.sh b/examples/benchmarks/distillation/submit_inference_benchmark.sh
deleted file mode 100755
index 6608e6852..000000000
--- a/examples/benchmarks/distillation/submit_inference_benchmark.sh
+++ /dev/null
@@ -1,45 +0,0 @@
-#!/bin/bash -l
-#SBATCH --partition=gpu-test
-#SBATCH --gres=gpu:80gb:1
-# Check the H100 GRES name with: sinfo -o "%P %G %f" (gpu:80gb may also match an A100)
-# Whole node, so other jobs do not distort the wall times
-#SBATCH --exclusive
-# Standard output and error:
-#SBATCH -o job_%j.out        # Standard output, %j = job ID
-#SBATCH -e job_%j.err        # Standard error, %j = job ID
-# Initial working directory (submit from this folder):
-#SBATCH -D ./
-# Job Name:
-#SBATCH -J distill_inference
-# Number of nodes
-#SBATCH --nodes=1
-#SBATCH --cpus-per-task=4
-
-# Time the energy and force inference of a teacher against its distilled students.
-# Set the paths below, then submit from this folder: sbatch submit_inference_benchmark.sh
-
-CONTAINER=/home/Jonas/container/schnetpack/schnetpack_container.sif
-OUTPUT=$SLURM_SUBMIT_DIR/results_$SLURM_JOB_ID
-
-# a MACE archive, from schnetpack.interfaces.mace_interface.export_mace_off
-TEACHER=/path/to/MACE-OFF23_medium.pt2
-# the best_model of each distillation run, as LABEL=PATH
-STUDENTS=(
-  student=/path/to/runs/<run_id>/best_model
-)
-# structures with the elements of the teacher, cycled to fill a batch
-STRUCTURES=/path/to/structures.xyz
-
-# pin the torch CPU threads to the allocated cores
-export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
-
-apptainer exec --nv $CONTAINER python $SLURM_SUBMIT_DIR/run_inference_benchmark.py \
-  --device cuda \
-  --teacher $TEACHER \
-  --teacher-format mace \
-  --students "${STUDENTS[@]}" \
-  --structures $STRUCTURES \
-  --batch-sizes 1 2 4 8 16 32 64 128 256 512 1024 2048 4096 \
-  --repeats 20 \
-  --warmup 5 \
-  --output $OUTPUT
diff --git a/src/schnetpack/dynamics/neighborlist.py b/src/schnetpack/dynamics/neighborlist.py
index 042f6505b..65a980192 100644
--- a/src/schnetpack/dynamics/neighborlist.py
+++ b/src/schnetpack/dynamics/neighborlist.py
@@ -1,4 +1,4 @@
-"""Neighbor lists for a batch of structures that keeps moving.
+"""Neighbor lists for a batch of structures.
 
 A relaxation or an MD run evaluates the same structures over and over, moving them a
 little each step. Rebuilding their neighbor lists every time is wasteful, so
@@ -29,19 +29,10 @@ class BatchNeighborList:
     """Keeps the neighbor lists of a batch valid while its structures move.
 
     The batch is a schnetpack input dictionary and stays on its device throughout, except
-    on the steps that have to rebuild: the neighbor list implementations run on cpu, so
-    the structures needing a new list -- and only those -- are shipped over and back.
+    on the steps that have to rebuild, since the the neighbor list implementations run on
+    cpu only.
 
-    ``idx_i``, ``idx_j`` and ``offsets`` are what this produces; the ``Rij`` that can come
-    with them is a by-product of the pruning, and models recompute it with their
-    ``PairwiseDistances`` input module anyway.
-
-    Note:
-        The offsets of a pair list are cell shifts, so pruning takes ``R[j] - R[i] +
-        offsets`` for the pair vector -- the convention ``PairwiseDistances`` and every
-        neighbor list in schnetpack share. Atoms that have wandered outside their cell
-        break it, here and in the model alike, so a structure with a cell has to stay in
-        it. Structures without one are unaffected.
+    Produces ``idx_i``, ``idx_j`` and ``offsets``.
 
     Args:
         neighbor_list: the neighbor list transform to build with, e.g.
@@ -267,13 +258,9 @@ def split_batch(
 
     The inverse of :func:`_atoms_collate_fn` for the structure-defining entries: atom-wise
     entries are cut along the atom axis at the ``n_atoms`` boundaries, structure-wise ones
-    are indexed. Every structure gets its position in the batch as ``properties.idx``, the
-    sample index a per-sample transform keys its caches by, so a batch that carries none --
-    one read back from a trajectory, say -- still splits into usable samples.
+    are indexed.
 
-    Neighbor lists are deliberately not split. They are the one thing that cannot be
-    recovered by cutting: the pair indices are shifted into batch-global numbering, and any
-    caller splitting a batch is about to rebuild them anyway.
+    Neighbor lists are deliberately not split, but need to be recalculated for each structure.
 
     Args:
         inputs: collated input batch.
diff --git a/src/schnetpack/interfaces/mace_interface.py b/src/schnetpack/interfaces/mace_interface.py
index e4b39621c..938afef86 100644
--- a/src/schnetpack/interfaces/mace_interface.py
+++ b/src/schnetpack/interfaces/mace_interface.py
@@ -239,6 +239,7 @@ def build_mace(
     radial_MLP: Sequence[int] = (64, 64, 64),
     MLP_irreps: str = "16x0e",
     scale: float | None = None,
+    atomic_numbers: Sequence[int] | None = None,
     seed: int = 0,
 ) -> str:
     """Build a MACE model with random weights for ``dataset``, export it (see
@@ -249,7 +250,8 @@ def build_mace(
     training data when it builds a model is taken from ``dataset``:
 
     - the elements, from all of ``dataset``, so that no validation or test
-      structure has an element the model does not cover;
+      structure has an element the model does not cover, unless
+      ``atomic_numbers`` is given;
     - the average number of neighbors within ``r_max``, the normalization of
       MACE's message passing, from ``indices``;
     - the scale of the atomic energies, as the RMS of the force components
@@ -262,6 +264,11 @@ def build_mace(
     train it with ``e0s="estimated"`` (see :class:`~schnetpack.model.MaceModel`).
     Any per-atom energy shift goes into the E0s, which are fitted per element.
 
+    ``atomic_numbers`` gives the model elements beyond those of ``dataset``,
+    e.g. those of a pretrained model to compare with: each element has weights
+    of its own, so the model has as many parameters as one built for these
+    elements. The weights of elements ``dataset`` lacks get no gradient.
+
     Args:
         path: where the archive is written.
         dataset: the data the model is trained on.
@@ -279,6 +286,8 @@ def build_mace(
         radial_MLP: hidden layer widths of the radial MLP.
         MLP_irreps: hidden irreps of the readout MLP.
         scale: scale of the atomic energies; None for the RMS of the forces.
+        atomic_numbers: the elements of the model, which must cover those of
+            ``dataset``; None for those of ``dataset``.
         seed: seed of the random weights.
     """
     with _mace_environment():
@@ -288,6 +297,13 @@ def build_mace(
         from mace import modules
 
         numbers = _elements(dataset)
+        if atomic_numbers is not None:
+            missing = sorted(set(numbers) - set(atomic_numbers))
+            if missing:
+                raise ValueError(
+                    f"atomic_numbers do not cover the elements {missing} of the dataset"
+                )
+            numbers = sorted(set(atomic_numbers))
         if scale is None:
             scale = _force_rms(dataset, indices, forces_key)
         config = {
diff --git a/src/schnetpack/transform/neighborlist.py b/src/schnetpack/transform/neighborlist.py
index a897d98aa..4d49f31b6 100644
--- a/src/schnetpack/transform/neighborlist.py
+++ b/src/schnetpack/transform/neighborlist.py
@@ -193,8 +193,14 @@ def _build_neighbor_list(self, Z, positions, cell, pbc, cutoff):
         # Check if shifts are needed for periodic boundary conditions
         if torch.all(pbc == 0):
             shifts = torch.zeros(0, 3, device=cell.device, dtype=torch.long)
+            wrap = torch.zeros_like(positions, dtype=torch.long)
         else:
             shifts = self._get_shifts(cell, pbc, cutoff)
+            # The shifts only reach far enough for atoms inside the cell, so search on
+            # wrapped positions and fold the wrapping back into the offsets below
+            wrap = torch.floor(positions.detach() @ torch.linalg.inv(cell)).long()
+            wrap = torch.where(pbc, wrap, 0)
+            positions = positions - wrap.to(cell.dtype) @ cell
         idx_i, idx_j, offset = self._get_neighbor_pairs(positions, cell, shifts, cutoff)
 
         # Create bidirectional id arrays, similar to what the ASE neighbor_list returns
@@ -207,7 +213,7 @@ def _build_neighbor_list(self, Z, positions, cell, pbc, cutoff):
         idx_j = bi_idx_j[sorted_idx]
 
         bi_offset = torch.cat((-offset, offset), dim=0)
-        offset = bi_offset[sorted_idx]
+        offset = bi_offset[sorted_idx] + wrap[idx_i] - wrap[idx_j]
         offset = torch.mm(offset.to(cell.dtype), cell)
 
         return idx_i, idx_j, offset
diff --git a/src/schnetpack/utils/compatibility.py b/src/schnetpack/utils/compatibility.py
index 629b2f4c0..80def92ed 100644
--- a/src/schnetpack/utils/compatibility.py
+++ b/src/schnetpack/utils/compatibility.py
@@ -53,6 +53,10 @@ def load_model(
     ``schnetpack.atomistic.*``, now below ``schnetpack.model``) are remapped before
     unpickling, see `_LEGACY_MODULES`.
 
+    Warning:
+        Model files are pickled PyTorch objects. Loading one executes code contained in the file, so only load
+        models from sources you trust.
+
     Args:
         model_path (str): Path to the saved model file.
         device (torch.device or str): Device on which to load the model. Defaults to "cpu".
diff --git a/tests/data/test_transforms.py b/tests/data/test_transforms.py
index 4b24935be..6b30aa832 100644
--- a/tests/data/test_transforms.py
+++ b/tests/data/test_transforms.py
@@ -51,6 +51,32 @@ def test_neighbor_list(self, neighbor_list, environment):
         for nbl, nbl_ref in zip(neighbors, neighbors_ref):
             torch.testing.assert_close(nbl, nbl_ref)
 
+    def test_atoms_outside_cell(self, neighbor_list, environment_periodic):
+        """Moving atoms by lattice vectors leaves the periodic system unchanged, so the
+        pair vectors must be the same as for the atoms inside the cell."""
+        cutoff, props, _ = environment_periodic
+        cell = props[structure.cell].view(3, 3)
+        lattice_shifts = torch.tensor(
+            [[1, 0, 0], [-2, 1, 0], [0, 0, 0], [1, -1, 3], [0, 2, -1]],
+            dtype=cell.dtype,
+        )
+        moved = {k: v.clone() for k, v in props.items()}
+        moved[structure.R] = props[structure.R] + lattice_shifts @ cell
+
+        pair_vectors = []
+        for inputs in (props, moved):
+            neighbors = neighbor_list(cutoff)(inputs)
+            R = inputs[structure.R]
+            neighbors[structure.Rij] = (
+                R[neighbors[structure.idx_j]]
+                - R[neighbors[structure.idx_i]]
+                + neighbors[structure.offsets]
+            )
+            pair_vectors.append(self._sort_neighbors(neighbors))
+
+        for nbl, nbl_ref in zip(*pair_vectors):
+            torch.testing.assert_close(nbl, nbl_ref)
+
     def _sort_neighbors(self, neighbors):
         """
         Routine for sorting the index, shift and distance vectors to allow comparison between different
diff --git a/tests/interfaces/test_mace_interface.py b/tests/interfaces/test_mace_interface.py
index e83614570..07ffa614b 100644
--- a/tests/interfaces/test_mace_interface.py
+++ b/tests/interfaces/test_mace_interface.py
@@ -196,6 +196,25 @@ def test_the_neighbors_are_counted_as_mace_counts_them(toy_data, r_max):
     )
 
 
+def test_atomic_numbers_give_the_model_more_elements(toy_data, tmp_path):
+    dataset, _ = toy_data
+    path = str(tmp_path / "mace.pt2")
+
+    core = load_mace(
+        build_mace(path, dataset, range(8), atomic_numbers=[8, 1, 7, 6], **SMALL)
+    )
+
+    assert core.atomic_numbers.tolist() == [1, 6, 7, 8]
+    assert not core.atomic_energies.any()
+
+
+def test_atomic_numbers_must_cover_the_dataset(toy_data, tmp_path):
+    dataset, _ = toy_data
+
+    with pytest.raises(ValueError, match=r"\[8\]"):
+        build_mace(str(tmp_path / "mace.pt2"), dataset, atomic_numbers=[1, 6])
+
+
 def test_the_seed_fixes_the_weights(built, toy_data, tmp_path):
     dataset, _ = toy_data
     first = load_mace(built).state_dict()