Hack better

metatensor · Jan 26, 2024 · addd65d · addd65d
1 parent 236e5a3
commit addd65d
Showing 1 changed file with 11 additions and 10 deletions.
diff --git 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],