Linting fix

src/metatrain/experimental/pet/model.py

@@ -19,7 +19,7 @@
 
 from ...utils.additive import ZBL
 from ...utils.dtype import dtype_to_str
-from .utils import systems_to_batch_dict, load_raw_pet_model
+from .utils import load_raw_pet_model, systems_to_batch_dict
 
 
 logger = logging.getLogger(__name__)

src/metatrain/experimental/pet/trainer.py

@@ -19,12 +19,7 @@
     update_pyg_graphs,
 )
 from pet.hypers import Hypers, save_hypers
-from pet.pet import (
-    PET,
-    FlagsWrapper,
-    PETMLIPWrapper,
-    PETUtilityWrapper,
-)
+from pet.pet import PET, FlagsWrapper, PETMLIPWrapper, PETUtilityWrapper
 from pet.utilities import (
     FullLogger,
     ModelKeeper,
@@ -44,7 +39,7 @@
 from ...utils.data import Dataset, check_datasets
 from ...utils.io import check_file_extension
 from . import PET as WrappedPET
-from .utils import dataset_to_ase, update_hypers, load_raw_pet_model
+from .utils import dataset_to_ase, load_raw_pet_model, update_hypers
 
 
 logger = logging.getLogger(__name__)

src/metatrain/experimental/pet/utils/load_raw_pet_model.py

@@ -1,17 +1,19 @@
-from pet.pet import PET, SelfContributionsWrapper
+from typing import Dict, List
+
+import numpy as np
 from pet.hypers import Hypers
+from pet.pet import PET, SelfContributionsWrapper
+
 from .fine_tuning import LoRAWrapper
-from . import update_state_dict
-from typing import List, Dict
-import numpy as np
+from .update_state_dict import update_state_dict
 
 
 def load_raw_pet_model(
     state_dict: Dict,
     hypers: Dict,
     atomic_types: List,
     self_contributions: np.ndarray,
-) -> PET:
+) -> "SelfContributionsWrapper":
     """Creates a raw PET model instance."""
 
     ARCHITECTURAL_HYPERS = Hypers(hypers)