fix imports formatting issues

ACEsuit · Feb 4, 2025 · bb2162c · bb2162c
1 parent 2da3ef6
commit bb2162c
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 14 deletions.
diff --git a/mace/__init__.py b/mace/__init__.py
@@ -1,5 +1,5 @@
-from .__version__ import __version__
-
 import os
 
+from .__version__ import __version__
+
 os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"
diff --git a/mace/calculators/mace.py b/mace/calculators/mace.py
@@ -3,15 +3,16 @@
 # Authors: Ilyes Batatia, David Kovacs
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
-import os
-
-os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"
 
+# pylint: disable=wrong-import-position
+import os
 import logging
 from glob import glob
 from pathlib import Path
 from typing import List, Union
 
+os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"
+
 import numpy as np
 import torch
 from ase.calculators.calculator import Calculator, all_changes

diff --git a/mace/cli/create_lammps_model.py b/mace/cli/create_lammps_model.py
@@ -1,9 +1,9 @@
+# pylint: disable=wrong-import-position
 import os
+import argparse
 
 os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"
 
-import argparse
-
 import torch
 from e3nn.util import jit
 

diff --git a/mace/data/utils.py b/mace/data/utils.py
@@ -271,9 +271,7 @@ def load_from_xyz(
             if isolated_atom_config:
                 atomic_number = int(atoms.get_atomic_numbers()[0])
                 if energy_key in atoms.info.keys():
-                    atomic_energies_dict[atomic_number] = float(atoms.info[
-                        energy_key
-                    ])
+                    atomic_energies_dict[atomic_number] = float(atoms.info[energy_key])
                 else:
                     logging.warning(
                         f"Configuration '{idx}' is marked as 'IsolatedAtom' "

diff --git a/tests/test_cueq.py b/tests/test_cueq.py
@@ -1,10 +1,10 @@
+# pylint: disable=wrong-import-position
 import os
-
-os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"
-
 from copy import deepcopy
 from typing import Any, Dict
 
+os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1"
+
 import pytest
 import torch
 import torch.nn.functional as F

diff --git a/tests/test_foundations.py b/tests/test_foundations.py
@@ -195,7 +195,7 @@ def test_compile_foundation(calc):
     model = calc.models[0]
     atoms = molecule("CH4")
     atoms.positions += np.random.randn(*atoms.positions.shape) * 0.1
-    batch = calc._atoms_to_batch(atoms)
+    batch = calc._atoms_to_batch(atoms)  # pylint disable=protected-access
     output_1 = model(batch.to_dict())
     model_compiled = jit.compile(model)
     output = model_compiled(batch.to_dict())