From addd65d941d10a798de6661efb37a129a7be2ac0 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Fri, 26 Jan 2024 14:33:18 +0100 Subject: [PATCH] Hack better --- .../soap_bpnn/tests/test_torchscript.py | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/src/metatensor/models/soap_bpnn/tests/test_torchscript.py b/src/metatensor/models/soap_bpnn/tests/test_torchscript.py index aaa20328c..4f5d0141b 100644 --- a/src/metatensor/models/soap_bpnn/tests/test_torchscript.py +++ b/src/metatensor/models/soap_bpnn/tests/test_torchscript.py @@ -1,9 +1,17 @@ import os -import torch -from metatensor.torch.atomistic import ModelCapabilities, ModelOutput -from metatensor.models.soap_bpnn import DEFAULT_HYPERS, Model +# Execute the setup script which will make sum_over_samples saveable. +current_dir = os.path.dirname(__file__) +setup_path = os.path.join( + current_dir, "..", "..", "..", "..", "..", "scripts", "setup.py" +) +exec(open(setup_path).read()) + +import torch # noqa: E402 +from metatensor.torch.atomistic import ModelCapabilities, ModelOutput # noqa: E402 + +from metatensor.models.soap_bpnn import DEFAULT_HYPERS, Model # noqa: E402 def test_torchscript(): @@ -26,13 +34,6 @@ def test_torchscript(): def test_torchscript_save(): """Tests that the model can be jitted and saved.""" - # Execute the setup script which will make sum_over_samples saveable. - current_dir = os.path.dirname(__file__) - setup_path = os.path.join( - current_dir, "..", "..", "..", "..", "..", "scripts", "setup.py" - ) - exec(open(setup_path).read()) - capabilities = ModelCapabilities( length_unit="Angstrom", species=[1, 6, 7, 8],