diff --git a/.github/workflows/linting_formatting.yml b/.github/workflows/linting_formatting.yml
new file mode 100644
index 0000000..7235de2
--- /dev/null
+++ b/.github/workflows/linting_formatting.yml
@@ -0,0 +1,27 @@
+name: Linting and Formatting
+
+on:
+  pull_request:
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v7
+        with:
+          persist-credentials: false
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          version: "0.12.0"
+
+      - name: Install package
+        run: uv sync --group dev
+
+      - name: Check styling
+        run: uv run ruff format --check .
+
+      - name: Check linting
+        run: uv run ruff check src/lanfactory
diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml
index 23e1d2a..f54bf6f 100644
--- a/.github/workflows/run_tests.yml
+++ b/.github/workflows/run_tests.yml
@@ -19,24 +19,20 @@ jobs:
         with:
           persist-credentials: false
 
-      - name: Setup Python
-        uses: actions/setup-python@v6
-        with:
-          python-version: ${{ matrix.python-version }}
-
       - name: Install uv
         uses: astral-sh/setup-uv@v7
         with:
-          version: "0.6.5"
+          version: "0.12.0"
+          python-version: ${{ matrix.python-version }}
           # enable-cache: true
           # cache-dependency-glob: "pyproject.toml pdm.lock"
 
       - name: Clear all caches
         run: |
-            rm -rf ~/.cache/pip
-            rm -rf ~/.cache/uv
-            rm -rf ~/.cache/conda
-            rm -rf ~/.cache/npm
+          rm -rf ~/.cache/pip
+          rm -rf ~/.cache/uv
+          rm -rf ~/.cache/conda
+          rm -rf ~/.cache/npm
 
       - name: Install package
         run: uv sync --all-groups --reinstall
@@ -47,12 +43,6 @@ jobs:
       - name: Run pytest
         run: uv run pytest
 
-      - name: Check styling
-        run: uv run ruff format --check .
-
-      - name: Linting
-        run: uv run ruff check src/lanfactory
-
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v7
         with:
@@ -67,15 +57,11 @@ jobs:
         with:
           persist-credentials: false
 
-      - name: Setup Python
-        uses: actions/setup-python@v6
-        with:
-          python-version: "3.12"
-
       - name: Install uv
         uses: astral-sh/setup-uv@v7
         with:
-          version: "0.6.5"
+          version: "0.12.0"
+          python-version: ${{ matrix.python-version }}
 
       - name: Install package (with notebook + backend deps)
         run: uv sync --all-groups
diff --git a/.gitignore b/.gitignore
index aa3bad6..df0292f 100755
--- a/.gitignore
+++ b/.gitignore
@@ -157,3 +157,4 @@ explorations/
 __marimo__/
 
 uv.lock
+.vscode/
diff --git a/pyproject.toml b/pyproject.toml
index 7277bf6..79f16d2 100755
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -82,7 +82,7 @@ dev = [
     "pytest-timer>=1.0.0",
     "pytest-xdist>=3.6.1",
     "pytest>=8.3.1",
-    "ruff>=0.14.4",
+    "ruff>=0.15.1",
     "types-PyYAML",
     "mlflow>=3.14.0",
     "jaxonnxruntime>=0.3",
@@ -120,6 +120,14 @@ line-length = 88
 src = ["src/lanfactory", "tests"]
 exclude = ["notebooks/*", "docs/*", "docs/basic_tutorial/basic_tutorial.ipynb"]
 
+[tool.ruff.lint]
+ignore = [
+    "B008",   # Function calls in default arguments (Typer/FastAPI options)
+    "EXE001", # Shebang present but file is not executable
+    "EXE002", # Executable file fails to declare a shebang
+    "BLE001", #
+]
+
 # Coverage settings
 [tool.coverage.run]
 source = ["src"]
diff --git a/src/lanfactory/__init__.py b/src/lanfactory/__init__.py
index 780a2f4..33b3b8f 100755
--- a/src/lanfactory/__init__.py
+++ b/src/lanfactory/__init__.py
@@ -1,11 +1,8 @@
 __version__ = "0.8.0"
 
-from . import config
-from . import trainers
-from . import utils
-from . import onnx
+from . import config, onnx, trainers, utils
 
-__all__ = ["config", "trainers", "utils", "onnx", "network_inspectors"]
+__all__ = ["config", "network_inspectors", "onnx", "trainers", "utils"]
 
 
 def __getattr__(name):
diff --git a/src/lanfactory/cli/download_hf.py b/src/lanfactory/cli/download_hf.py
index 262f776..7001b45 100644
--- a/src/lanfactory/cli/download_hf.py
+++ b/src/lanfactory/cli/download_hf.py
@@ -12,7 +12,6 @@
 from pathlib import Path
 
 import typer
-
 from lanfactory.hf import DEFAULT_REPO_ID, VALID_NETWORK_TYPES
 
 app = typer.Typer()
diff --git a/src/lanfactory/cli/jax_train.py b/src/lanfactory/cli/jax_train.py
index ffde030..0132ef7 100644
--- a/src/lanfactory/cli/jax_train.py
+++ b/src/lanfactory/cli/jax_train.py
@@ -164,9 +164,10 @@ def main(
 
     if mlflow_tracking_enabled:
         try:
-            import mlflow
             import os
 
+            import mlflow
+
             # Set tracking URI with priority: CLI arg > env var > default
             if mlflow_tracking_uri:
                 tracking_uri = mlflow_tracking_uri
@@ -234,9 +235,10 @@ def main(
 
             if not mlflow_tracking_enabled:
                 # Need to initialize MLflow just for querying
-                import mlflow
                 import os
 
+                import mlflow
+
                 # Use same logic as above for tracking URI
                 if mlflow_tracking_uri:
                     tracking_uri = mlflow_tracking_uri
@@ -312,7 +314,7 @@ def main(
     # Mode 2: Validation - verify MLflow files exist in training_data_folder
     if mlflow_lineage_info and training_data_folder:
         expected_files = set(mlflow_lineage_info["all_files"])
-        actual_files = set(f.name for f in valid_file_list)
+        actual_files = {f.name for f in valid_file_list}
 
         missing_files = expected_files - actual_files
         extra_files = actual_files - expected_files
@@ -463,13 +465,8 @@ def main(
         ]
     )
 
-    pickle.dump(
-        network_config,
-        open(
-            networks_path / file_name_suffix,
-            "wb",
-        ),
-    )
+    file_path = networks_path / file_name_suffix
+    file_path.write_bytes(pickle.dumps(network_config))
 
     # Load network
     net = lanfactory.trainers.JaxMLPFactory(
diff --git a/src/lanfactory/cli/torch_train.py b/src/lanfactory/cli/torch_train.py
index 84b1855..c82bfb6 100644
--- a/src/lanfactory/cli/torch_train.py
+++ b/src/lanfactory/cli/torch_train.py
@@ -22,11 +22,10 @@
 from importlib.resources import as_file, files
 from pathlib import Path
 
+import lanfactory
 import psutil
 import torch
 import typer
-
-import lanfactory
 from lanfactory.cli.utils import (
     _get_train_network_config,
 )
@@ -168,9 +167,10 @@ def main(
 
     if mlflow_tracking_enabled:
         try:
-            import mlflow
             import os
 
+            import mlflow
+
             # Set tracking URI with priority: CLI arg > env var > default
             if mlflow_tracking_uri:
                 tracking_uri = mlflow_tracking_uri
@@ -238,9 +238,10 @@ def main(
 
             if not mlflow_tracking_enabled:
                 # Need to initialize MLflow just for querying
-                import mlflow
                 import os
 
+                import mlflow
+
                 # Use same logic as above for tracking URI
                 if mlflow_tracking_uri:
                     tracking_uri = mlflow_tracking_uri
@@ -316,7 +317,7 @@ def main(
     # Mode 2: Validation - verify MLflow files exist in training_data_folder
     if mlflow_lineage_info and training_data_folder:
         expected_files = set(mlflow_lineage_info["all_files"])
-        actual_files = set(f.name for f in valid_file_list)
+        actual_files = {f.name for f in valid_file_list}
 
         missing_files = expected_files - actual_files
         extra_files = actual_files - expected_files
@@ -467,13 +468,8 @@ def main(
         ]
     )
 
-    pickle.dump(
-        network_config,
-        open(
-            networks_path / file_name_suffix,
-            "wb",
-        ),
-    )
+    file_path = networks_path / file_name_suffix
+    file_path.write_bytes(pickle.dumps(network_config))
 
     # Load network
     net = lanfactory.trainers.TorchMLP(
diff --git a/src/lanfactory/cli/upload_hf.py b/src/lanfactory/cli/upload_hf.py
index 538b1db..4142a83 100644
--- a/src/lanfactory/cli/upload_hf.py
+++ b/src/lanfactory/cli/upload_hf.py
@@ -12,7 +12,6 @@
 from pathlib import Path
 
 import typer
-
 from lanfactory.hf import DEFAULT_REPO_ID, VALID_NETWORK_TYPES
 
 app = typer.Typer()
diff --git a/src/lanfactory/cli/utils.py b/src/lanfactory/cli/utils.py
index 7933973..dc13f6b 100644
--- a/src/lanfactory/cli/utils.py
+++ b/src/lanfactory/cli/utils.py
@@ -1,10 +1,11 @@
 # import argparse
 import logging
-from pathlib import Path
 import pickle
-import yaml
-import numpy as np
+from pathlib import Path
+
 import lanfactory
+import numpy as np
+import yaml
 
 logger = logging.getLogger(__name__)
 
@@ -88,9 +89,9 @@ def _make_train_network_configs(
     if save_name:
         save_folder = Path(save_folder)
         save_folder.mkdir(parents=True, exist_ok=True)  # pragma: no cover
-        save_name = save_folder / save_name
-        pickle.dump(config_dict, open(save_name, "wb"))
-        print(f"Saved to: {save_name}")
+        save_path = save_folder / save_name
+        save_path.write_bytes(pickle.dumps(config_dict))
+        print(f"Saved to: {save_path}")
     else:
         print("No save name provided, config not saved to file.")
 
@@ -99,7 +100,7 @@ def _make_train_network_configs(
 
 def _get_train_network_config(yaml_config_path: str | Path | None = None, net_index=0):
     if yaml_config_path is not None:
-        basic_config = yaml.safe_load(open(yaml_config_path, "rb"))
+        basic_config = yaml.safe_load(Path(yaml_config_path).read_bytes())
         network_type = basic_config["NETWORK_TYPE"]
     else:
         raise ValueError("No YAML config path provided")
diff --git a/src/lanfactory/config/__init__.py b/src/lanfactory/config/__init__.py
index 1878766..42aee52 100755
--- a/src/lanfactory/config/__init__.py
+++ b/src/lanfactory/config/__init__.py
@@ -1,21 +1,21 @@
 from .network_configs import (
-    network_config_mlp,
     network_config_choice_prob,
-    network_config_opn,
     network_config_cpn,
-    train_config_mlp,
+    network_config_mlp,
+    network_config_opn,
     train_config_choice_prob,
-    train_config_opn,
     train_config_cpn,
+    train_config_mlp,
+    train_config_opn,
 )
 
 __all__ = [
-    "network_config_mlp",
     "network_config_choice_prob",
-    "network_config_opn",
     "network_config_cpn",
-    "train_config_mlp",
+    "network_config_mlp",
+    "network_config_opn",
     "train_config_choice_prob",
-    "train_config_opn",
     "train_config_cpn",
+    "train_config_mlp",
+    "train_config_opn",
 ]
diff --git a/src/lanfactory/hf/__init__.py b/src/lanfactory/hf/__init__.py
index 7dbda0a..f83e6c5 100644
--- a/src/lanfactory/hf/__init__.py
+++ b/src/lanfactory/hf/__init__.py
@@ -7,20 +7,20 @@
 DEFAULT_REPO_ID = "franklab/HSSM"
 VALID_NETWORK_TYPES = ("lan", "cpn", "opn")
 
-from lanfactory.hf.model_card import (  # noqa: E402
-    load_model_card_yaml,
-    generate_readme,
+from lanfactory.hf.download import download_model
+from lanfactory.hf.model_card import (
     ModelCardConfig,
+    generate_readme,
+    load_model_card_yaml,
 )
-from lanfactory.hf.upload import upload_model  # noqa: E402
-from lanfactory.hf.download import download_model  # noqa: E402
+from lanfactory.hf.upload import upload_model
 
 __all__ = [
     "DEFAULT_REPO_ID",
     "VALID_NETWORK_TYPES",
-    "load_model_card_yaml",
-    "generate_readme",
     "ModelCardConfig",
-    "upload_model",
     "download_model",
+    "generate_readme",
+    "load_model_card_yaml",
+    "upload_model",
 ]
diff --git a/src/lanfactory/hf/download.py b/src/lanfactory/hf/download.py
index 52bb17c..ecef4e7 100644
--- a/src/lanfactory/hf/download.py
+++ b/src/lanfactory/hf/download.py
@@ -123,9 +123,10 @@ def _download_model_hf(  # pragma: no cover
     model_files = [f for f in all_files if f.startswith(path_prefix)]
 
     if not model_files:
+        top_level_dirs = {f.split("/")[0] for f in all_files if "/" in f}
         raise FileNotFoundError(
             f"No files found at {repo_id}/{path_prefix}. "
-            f"Available paths: {set(f.split('/')[0] for f in all_files if '/' in f)}"
+            f"Available paths: {top_level_dirs}"
         )
 
     if include_patterns:
diff --git a/src/lanfactory/hf/upload.py b/src/lanfactory/hf/upload.py
index c223ef2..adbf23f 100644
--- a/src/lanfactory/hf/upload.py
+++ b/src/lanfactory/hf/upload.py
@@ -155,7 +155,8 @@ def _upload_to_hf(  # pragma: no cover
 ) -> str:
     """HF-dependent implementation of upload_model."""
     try:
-        from huggingface_hub import HfApi, create_repo as hf_create_repo
+        from huggingface_hub import HfApi
+        from huggingface_hub import create_repo as hf_create_repo
     except ImportError as exc:
         raise ImportError(
             "huggingface_hub is required for HuggingFace uploads. "
diff --git a/src/lanfactory/network_inspectors/__init__.py b/src/lanfactory/network_inspectors/__init__.py
index 6f8e2ab..6154967 100644
--- a/src/lanfactory/network_inspectors/__init__.py
+++ b/src/lanfactory/network_inspectors/__init__.py
@@ -15,10 +15,10 @@
 from .loaders import get_torch_mlp
 
 __all__ = [
+    "GridSpec",
+    "ModelSpec",
+    "PlotConfig",
     "get_torch_mlp",
     "kde_vs_lan_likelihoods",
     "lan_manifold",
-    "ModelSpec",
-    "PlotConfig",
-    "GridSpec",
 ]
diff --git a/src/lanfactory/network_inspectors/api.py b/src/lanfactory/network_inspectors/api.py
index 0635786..570cdbe 100644
--- a/src/lanfactory/network_inspectors/api.py
+++ b/src/lanfactory/network_inspectors/api.py
@@ -2,8 +2,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
 import logging
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
diff --git a/src/lanfactory/network_inspectors/config.py b/src/lanfactory/network_inspectors/config.py
index 1eee71c..2c621e6 100644
--- a/src/lanfactory/network_inspectors/config.py
+++ b/src/lanfactory/network_inspectors/config.py
@@ -8,7 +8,6 @@
 
 import numpy as np
 from numpy.typing import NDArray
-
 from ssms.config import ModelConfigBuilder
 
 
diff --git a/src/lanfactory/network_inspectors/loaders.py b/src/lanfactory/network_inspectors/loaders.py
index 4940051..88fa3a5 100644
--- a/src/lanfactory/network_inspectors/loaders.py
+++ b/src/lanfactory/network_inspectors/loaders.py
@@ -2,11 +2,11 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
 import os
+import warnings
+from collections.abc import Callable
 from os import PathLike
 from typing import Any
-import warnings
 
 import numpy as np
 from numpy.typing import NDArray
diff --git a/src/lanfactory/onnx/__init__.py b/src/lanfactory/onnx/__init__.py
index 0b25cb2..289ac30 100755
--- a/src/lanfactory/onnx/__init__.py
+++ b/src/lanfactory/onnx/__init__.py
@@ -3,7 +3,7 @@
 from .transform_onnx import transform_to_onnx
 
 __all__ = [
-    "transform_to_onnx",
-    "transform_sbi_to_onnx",
     "transform_bayesflow_to_onnx",
+    "transform_sbi_to_onnx",
+    "transform_to_onnx",
 ]
diff --git a/src/lanfactory/trainers/__init__.py b/src/lanfactory/trainers/__init__.py
index ecd495d..50627b4 100755
--- a/src/lanfactory/trainers/__init__.py
+++ b/src/lanfactory/trainers/__init__.py
@@ -1,32 +1,32 @@
 import warnings
 
+from .jax_mlp import JaxMLP, JaxMLPFactory, ModelTrainerJaxMLP
 from .torch_mlp import (
     DatasetTorch,
-    TorchMLP,
-    TorchMLPFactory,
-    ModelTrainerTorchMLP,
     LoadTorchMLP,
     LoadTorchMLPInfer,
+    ModelTrainerTorchMLP,
+    TorchMLP,
+    TorchMLPFactory,
     make_dataloader,
     make_train_valid_dataloaders,
 )
-from .jax_mlp import JaxMLPFactory, JaxMLP, ModelTrainerJaxMLP
 
 __all__ = [
     # Dataset and DataLoader helpers
     "DatasetTorch",
-    "make_dataloader",
-    "make_train_valid_dataloaders",
-    # Torch MLP
-    "TorchMLP",
-    "TorchMLPFactory",
-    "ModelTrainerTorchMLP",
-    "LoadTorchMLP",
-    "LoadTorchMLPInfer",
+    "JaxMLP",
     # Jax MLP
     "JaxMLPFactory",
-    "JaxMLP",
+    "LoadTorchMLP",
+    "LoadTorchMLPInfer",
     "ModelTrainerJaxMLP",
+    "ModelTrainerTorchMLP",
+    # Torch MLP
+    "TorchMLP",
+    "TorchMLPFactory",
+    "make_dataloader",
+    "make_train_valid_dataloaders",
 ]
 
 _DEPRECATED_ALIASES = {
diff --git a/src/lanfactory/trainers/jax_mlp.py b/src/lanfactory/trainers/jax_mlp.py
index 3c3968c..bf508f5 100755
--- a/src/lanfactory/trainers/jax_mlp.py
+++ b/src/lanfactory/trainers/jax_mlp.py
@@ -3,10 +3,11 @@
 """
 
 import pickle
+from collections.abc import Callable, Sequence
 from functools import partial
 from pathlib import Path
 from time import time
-from typing import Any, Callable, Sequence
+from typing import Any
 
 import flax
 import jax
@@ -49,7 +50,7 @@ def JaxMLPFactory(
     elif isinstance(network_config, dict):
         network_config_internal = network_config
     else:
-        raise ValueError(
+        raise TypeError(
             "network_config argument is not passed as either a dictionary or a string (path to a file)!"
         )
 
@@ -85,10 +86,6 @@ class JaxMLP(nn.Module):
     activations_dict = frozendict(
         {"relu": nn.relu, "tanh": nn.tanh, "sigmoid": nn.sigmoid}
     )
-    # network_type: Optional[str] = "none"
-
-    # Define network type
-    # network_type = "lan" if train_output_type == "logprob" else "cpn"
 
     def setup(self) -> None:
         """Setup function for the JaxMLP class.
@@ -123,20 +120,11 @@ def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
 
         for i, lyr in enumerate(self.layers):
             x = lyr(x)
-            if i != (len(self.layers) - 1):
+            if i != (len(self.layers) - 1) or self.activations[i] != "linear":
                 x = self.activation_funs[i](x)
-            else:
-                if self.activations[i] == "linear":
-                    pass
-                else:
-                    x = self.activation_funs[i](x)
-
-        if (not self.train) and (self.train_output_type == "logprob"):
-            x = x  # just for pedagogy
-        elif (not self.train) and (self.train_output_type == "logits"):
-            x = -jnp.log((1 + jnp.exp(-x)))
-        elif not self.train:  # pragma: no cover
-            x = x  # just for pedagogy
+
+        if (not self.train) and (self.train_output_type == "logits"):
+            x = -jnp.log(1 + jnp.exp(-x))
 
         return x
 
@@ -220,7 +208,7 @@ def make_forward_partial(
         elif isinstance(state, dict):
             loaded_state = state
         else:
-            raise ValueError("state argument has to be a dictionary or a string!")
+            raise TypeError("state argument has to be a dictionary or a string!")
 
         # Make forward pass
         net_forward = partial(self.apply, loaded_state)
@@ -272,7 +260,7 @@ def __init__(
                 The ModelTrainerJaxMLP object.
 
         """
-        if "loss_dict" not in train_config.keys():
+        if "loss_dict" not in train_config:
             self.loss_dict: dict[str, dict] = {
                 "huber": {"fun": optax.huber_loss, "kwargs": {"delta": 1}},
                 "mse": {"fun": optax.l2_loss, "kwargs": {}},
@@ -281,7 +269,7 @@ def __init__(
         else:  # pragma: no cover
             self.loss_dict = train_config["loss_dict"]
 
-        if "lr_dict" not in train_config.keys():
+        if "lr_dict" not in train_config:
             # Todo: Add more schedules (for now warmup_cosine_decay_schedule)
             self.lr_dict: dict[str, float] = {
                 "init_value": 0.0002,
@@ -335,10 +323,10 @@ def loss_fn(params: dict) -> tuple[float, jnp.ndarray]:
 
             if train:
                 grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
-                (loss, pred), grads = grad_fn(state.params)
+                (loss, _), grads = grad_fn(state.params)
                 return grads, loss
             else:
-                loss, pred = loss_fn(state.params)
+                loss, _ = loss_fn(state.params)
                 return loss
 
         return apply_model_core
@@ -431,7 +419,7 @@ def run_epoch(
         # Run training for one epoch
         start_time = time()
         step = 0
-        for X, y in tmp_dataloader:
+        for step, (X, y) in enumerate(tmp_dataloader):
             X_jax = jnp.array(X)
             y_jax = jnp.array(y)
 
@@ -459,24 +447,19 @@ def run_epoch(
                 if self.mlflow_on:
                     try:
                         mlflow.log_metric("loss", float(loss), step=int(state.step))
-                    except Exception:
-                        pass
-
-                elif verbose == 1:  # pragma: no cover
-                    if (step % 1000) == 0:
-                        print(
-                            train_str
-                            + " - Step: "
-                            + str(step)
-                            + " of "
-                            + str(cnt_max)
-                            + " - Loss: "
-                            + str(loss)
-                        )
-                else:
-                    pass
-
-            step += 1
+                    except Exception as e:
+                        print(f"Failed to log metric to MLflow: {e}")
+
+                elif verbose == 1 and (step % 1000) == 0:  # pragma: no cover
+                    print(
+                        train_str
+                        + " - Step: "
+                        + str(step)
+                        + " of "
+                        + str(cnt_max)
+                        + " - Loss: "
+                        + str(loss)
+                    )
 
         end_time = time()
         print(
@@ -548,7 +531,7 @@ def train_and_evaluate(
 
         # Initialize network
         if not isinstance(self.seed, int):
-            raise ValueError(
+            raise TypeError(
                 "seed argument is not an integer, "
                 + "please specify a valid seed to make this code reproducible!"
             )
@@ -561,7 +544,7 @@ def train_and_evaluate(
         # Training loop over epochs
         for epoch in range(self.train_config["n_epochs"]):
             print("Epoch: " + str(epoch) + " of " + str(self.train_config["n_epochs"]))
-            state, train_loss = self.run_epoch(
+            state, _ = self.run_epoch(
                 state,
                 train=True,
                 verbose=verbose,
@@ -605,26 +588,23 @@ def train_and_evaluate(
 
             # Write to file
             train_state_path = f"{full_path}_train_state.jax"
-            file = open(train_state_path, "wb")
-            file.write(byte_output)
-            file.close()
+            Path(train_state_path).write_bytes(byte_output)
             print("Saving model parameters to: " + train_state_path)
 
-            config_path = f"{full_path}_train_config.pickle"
-            pickle.dump(self.train_config, open(config_path, "wb"))
-            print("Saving training config to: " + config_path)
-
-            data_details_path = f"{full_path}_data_details.pickle"
-            pickle.dump(
-                {
-                    "train_data_generator_config": self.train_dl.dataset.data_generator_config,
-                    "train_data_file_ids": self.train_dl.dataset.file_ids,
-                    "valid_data_generator_config": self.valid_dl.dataset.data_generator_config,
-                    "valid_data_file_ids": self.valid_dl.dataset.file_ids,
-                },
-                open(data_details_path, "wb"),
-            )
-            print("Saving training data details to: " + data_details_path)
+            config_path = Path(f"{full_path}_train_config.pickle")
+            config_path.write_bytes(pickle.dumps(self.train_config))
+            print(f"Saving training config to: {config_path}")
+
+            data_details_path = Path(f"{full_path}_data_details.pickle")
+            data_details = {
+                "train_data_generator_config": self.train_dl.dataset.data_generator_config,
+                "train_data_file_ids": self.train_dl.dataset.file_ids,
+                "valid_data_generator_config": self.valid_dl.dataset.data_generator_config,
+                "valid_data_file_ids": self.valid_dl.dataset.file_ids,
+            }
+
+            data_details_path.write_bytes(pickle.dumps(data_details))
+            print(f"Saving training data details to: {data_details_path}")
 
             if self.mlflow_on:
                 try:
diff --git a/src/lanfactory/trainers/torch_mlp.py b/src/lanfactory/trainers/torch_mlp.py
index d439bb7..8fc9708 100755
--- a/src/lanfactory/trainers/torch_mlp.py
+++ b/src/lanfactory/trainers/torch_mlp.py
@@ -1,18 +1,16 @@
 """This module contains the classes for training TorchMLP models."""
 
-import numpy as np
-import pandas as pd
-import pickle
-from typing import Callable
-from time import time
 import logging
+import pickle
+from collections.abc import Callable
 from pathlib import Path
+from time import time
 
-
+import numpy as np
+import pandas as pd
 import torch
-import torch.nn as nn
-import torch.optim as optim
 import torch.nn.functional as F
+from torch import nn, optim
 from torch.utils.data import DataLoader
 
 try:
@@ -90,7 +88,8 @@ def __getitem__(self, index: int) -> tuple[np.ndarray, np.ndarray]:
 
     def __load_file(self, file_index: int) -> None:
         # Load file and shuffle the indices
-        self.tmp_data = pickle.load(open(self.file_ids[file_index], "rb"))
+        with open(self.file_ids[file_index], "rb") as f:
+            self.tmp_data = pickle.load(f)
         shuffle_idx = np.random.choice(
             self.tmp_data[self.features_key].shape[0],
             size=self.tmp_data[self.features_key].shape[0],
@@ -100,11 +99,11 @@ def __load_file(self, file_index: int) -> None:
             shuffle_idx, :
         ]
         self.tmp_data[self.label_key] = self.tmp_data[self.label_key][shuffle_idx]
-        return
 
     def __init_file_shape(self) -> None:
         # Function gets dimensionalities form a test data file
-        init_file = pickle.load(open(self.file_ids[0], "rb"))
+        with open(self.file_ids[0], "rb") as f:
+            init_file = pickle.load(f)
         self.file_shape_dict = {
             "inputs": init_file[self.features_key].shape,
             "labels": init_file[self.label_key].shape,
@@ -130,7 +129,6 @@ def __init_file_shape(self) -> None:
             self.label_dim = self.file_shape_dict["labels"][1]
         else:
             self.label_dim = 1
-        return
 
     def __data_generation(
         self, batch_ids: np.ndarray | None = None
@@ -337,7 +335,7 @@ def TorchMLPFactory(
     """
     if isinstance(network_config, str):
         with open(network_config, "rb") as f:
-            network_config = pickle.load(f)  # noqa: S301
+            network_config = pickle.load(f)
 
     assert isinstance(network_config, dict)
     return TorchMLP(
@@ -366,7 +364,7 @@ def __init__(
         input_shape: int = 10,
         network_type: str | None = None,
     ) -> None:
-        super(TorchMLP, self).__init__()
+        super().__init__()
 
         self.input_shape = input_shape
         self.network_config = network_config
@@ -452,7 +450,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             return self.layers[-1](x)
         elif self.train_output_type == "logits":
             return -torch.log(
-                (1 + torch.exp(-self.layers[-1](x)))
+                1 + torch.exp(-self.layers[-1](x))
             )  # log ( 1 / (1 + exp(-x))), where x = log(p / (1 - p))
         else:
             return self.layers[-1](x)
@@ -497,11 +495,15 @@ def __init__(
         elif isinstance(train_config, str | Path):
             print("train_config is passed as string or path: \n", train_config)
             try:
-                logger.info("Trying to load string as path to pickle file: ")
-                self.train_config: dict = pickle.load(open(train_config, "rb"))
+                logger.info(
+                    "Trying to load string as path to pickle file: %s", train_config
+                )
+                self.train_config: dict = pickle.loads(Path(train_config).read_bytes())
             except (OSError, pickle.PickleError) as e:  # pragma: no cover
                 logger.error(
-                    f"Error loading training config from file {train_config}: {str(e)}"
+                    "Error loading training config from file %s: %s",
+                    train_config,
+                    e,
                 )
                 raise
         elif isinstance(train_config, dict):
@@ -578,36 +580,24 @@ def __get_scheduler(self) -> None:
                     self.optimizer,
                     mode="min",
                     factor=(
-                        self.train_config["lr_scheduler_params"]["factor"]
-                        if "factor" in self.train_config["lr_scheduler_params"]
-                        else 0.1
+                        self.train_config["lr_scheduler_params"].get("factor", 0.1)
                     ),
                     patience=(
-                        self.train_config["lr_scheduler_params"]["patience"]
-                        if "patience" in self.train_config["lr_scheduler_params"]
-                        else 2
+                        self.train_config["lr_scheduler_params"].get("patience", 2)
                     ),
                     threshold=(
-                        self.train_config["lr_scheduler_params"]["threshold"]
-                        if "threshold" in self.train_config["lr_scheduler_params"]
-                        else 0.001
+                        self.train_config["lr_scheduler_params"].get("threshold", 0.001)
                     ),
                     threshold_mode="rel",
                     cooldown=0,
                     min_lr=(
-                        self.train_config["lr_scheduler_params"]["min_lr"]
-                        if "min_lr" in self.train_config["lr_scheduler_params"]
-                        else 0.00000001
+                        self.train_config["lr_scheduler_params"].get("min_lr", 1e-08)
                     ),
                 )
             elif self.train_config["lr_scheduler"] == "multiply":
                 self.scheduler = optim.lr_scheduler.ExponentialLR(
                     self.optimizer,
-                    gamma=(
-                        self.train_config["lr_scheduler_params"]["factor"]
-                        if "factor" in self.train_config["lr_scheduler_params"]
-                        else 0.1
-                    ),
+                    gamma=(self.train_config["lr_scheduler_params"].get("factor", 0.1)),
                     last_epoch=-1,
                 )
 
@@ -665,7 +655,7 @@ def train_and_evaluate(
             epoch_s_t = time()
 
             # Training loop
-            for xb, yb in self.train_dl:
+            for cnt, (xb, yb) in enumerate(self.train_dl):
                 # Shift data to device
                 if self.pin_memory and str(self.dev) == "cuda":
                     xb, yb = xb.cuda(non_blocking=True), yb.cuda(non_blocking=True)
@@ -684,7 +674,6 @@ def train_and_evaluate(
                 # Log training progress
                 self._log_training_progress(epoch, cnt, loss, verbose)
 
-                cnt += 1
                 step_cnt += 1
 
             print(
diff --git a/src/lanfactory/utils/__init__.py b/src/lanfactory/utils/__init__.py
index 0e9effc..029cb29 100755
--- a/src/lanfactory/utils/__init__.py
+++ b/src/lanfactory/utils/__init__.py
@@ -1,11 +1,11 @@
-from .util_funs import save_configs
 from .mlflow_utils import (
     get_files_from_data_generation_experiment,
     log_training_data_lineage,
 )
+from .util_funs import save_configs
 
 __all__ = [
-    "save_configs",
     "get_files_from_data_generation_experiment",
     "log_training_data_lineage",
+    "save_configs",
 ]
diff --git a/src/lanfactory/utils/mlflow_utils.py b/src/lanfactory/utils/mlflow_utils.py
index 2d93193..d9de6bf 100644
--- a/src/lanfactory/utils/mlflow_utils.py
+++ b/src/lanfactory/utils/mlflow_utils.py
@@ -110,7 +110,7 @@ def log_training_data_lineage(
     training_data_folder: Path,
     valid_file_list: list,
     n_training_files: int,
-    tracking_uri: str = None,
+    tracking_uri: str | None = None,
 ) -> dict:
     """Log training data lineage information to MLflow.
 
@@ -141,8 +141,9 @@ def log_training_data_lineage(
             - extra_files
     """
     try:
-        import mlflow
         import os
+
+        import mlflow
     except ImportError:
         logger.error("mlflow package not installed")
         return {}
@@ -165,7 +166,7 @@ def log_training_data_lineage(
 
         # Verify we have all the files we expect
         expected_file_names = set(expected_files_info["all_files"])
-        actual_file_names = set(f.name for f in valid_file_list)
+        actual_file_names = {f.name for f in valid_file_list}
 
         missing_files = expected_file_names - actual_file_names
         extra_files = actual_file_names - expected_file_names
diff --git a/src/lanfactory/utils/util_funs.py b/src/lanfactory/utils/util_funs.py
old mode 100755
new mode 100644
index 6718380..589a5df
--- a/src/lanfactory/utils/util_funs.py
+++ b/src/lanfactory/utils/util_funs.py
@@ -28,14 +28,13 @@ def save_configs(
     Path(save_folder).mkdir(parents=True, exist_ok=True)
 
     # Save network config
-    pickle.dump(
-        network_config,
-        open(Path(save_folder) / f"{model_id}_network_config.pickle", "wb"),
-    )
+    network_config_path = Path(save_folder) / f"{model_id}_network_config.pickle"
+    with open(network_config_path, "wb") as f:
+        pickle.dump(network_config, f)
     print("Saved network config")
+
     # Save train config
-    pickle.dump(
-        train_config, open(Path(save_folder) / f"{model_id}_train_config.pickle", "wb")
-    )
+    train_config_path = Path(save_folder) / f"{model_id}_train_config.pickle"
+    with open(train_config_path, "wb") as f:
+        pickle.dump(train_config, f)
     print("Saved train config")
-    return
diff --git a/tests/conftest.py b/tests/conftest.py
index ae0ae3e..f99b1bd 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,6 +1,5 @@
 import logging
 import multiprocessing
-import os
 import random
 import time
 import uuid
@@ -133,7 +132,7 @@ def _select_model(
 
 
 @pytest.fixture
-def dummy_generator_config(model_selector):
+def dummy_generator_config(model_selector, tmp_path):
     """Fixture providing a dummy model config for testing."""
 
     def _dummy_generator_config(mode="random"):
@@ -148,9 +147,7 @@ def _dummy_generator_config(mode="random"):
             generator_config["simulator"]["n_samples"] = (
                 TEST_GENERATOR_CONSTANTS.N_SAMPLES
             )
-            generator_config["output"]["folder"] = os.path.join(
-                TEST_GENERATOR_CONSTANTS.OUT_FOLDER, str(uuid.uuid4())
-            )
+            generator_config["output"]["folder"] = str(tmp_path / str(uuid.uuid4()))
             generator_config["training"]["n_samples_per_param"] = (
                 TEST_GENERATOR_CONSTANTS.N_SAMPLES_BY_PARAMETER_SET
             )
@@ -169,7 +166,7 @@ def _dummy_generator_config(mode="random"):
 
 
 @pytest.fixture
-def dummy_generator_config_simple_two_choices(model_selector):
+def dummy_generator_config_simple_two_choices(model_selector, tmp_path):
     """Fixture providing a dummy model config for testing."""
 
     def _dummy_generator_config_simple_two_choices(mode="random"):
@@ -185,9 +182,8 @@ def _dummy_generator_config_simple_two_choices(mode="random"):
             generator_config["simulator"]["n_samples"] = (
                 TEST_GENERATOR_CONSTANTS.N_SAMPLES
             )
-            generator_config["output"]["folder"] = os.path.join(
-                TEST_GENERATOR_CONSTANTS.OUT_FOLDER, str(uuid.uuid4())
-            )
+            generator_config["output"]["folder"] = str(tmp_path / str(uuid.uuid4()))
+
             generator_config["training"]["n_samples_per_param"] = (
                 TEST_GENERATOR_CONSTANTS.N_SAMPLES_BY_PARAMETER_SET
             )
diff --git a/tests/constants.py b/tests/constants.py
index dd65947..6404efa 100644
--- a/tests/constants.py
+++ b/tests/constants.py
@@ -31,7 +31,6 @@ class TestGeneratorConstants:
     N_TRAINING_SAMPLES: int = 2000
     N_SAMPLES_BY_PARAMETER_SET: int = 2000
     TEST_FOLDER: str = "tests/test_data"
-    OUT_FOLDER: str = "tests/test_data/lan_mlp/training_data"
     N_DATA_FILES: int = 2
     DEVICE: str = "cpu"
 
diff --git a/tests/hf/test_download.py b/tests/hf/test_download.py
index 686c185..47bee5b 100644
--- a/tests/hf/test_download.py
+++ b/tests/hf/test_download.py
@@ -3,7 +3,6 @@
 from unittest.mock import patch
 
 import pytest
-
 from lanfactory.hf.download import (
     DEFAULT_REPO_ID,
     download_model,
diff --git a/tests/hf/test_model_card.py b/tests/hf/test_model_card.py
index 82e3525..c9a02a4 100644
--- a/tests/hf/test_model_card.py
+++ b/tests/hf/test_model_card.py
@@ -4,7 +4,6 @@
 
 import pytest
 import yaml
-
 from lanfactory.hf.model_card import (
     ModelCardConfig,
     generate_readme,
diff --git a/tests/hf/test_upload.py b/tests/hf/test_upload.py
index 6eea1f6..fda6272 100644
--- a/tests/hf/test_upload.py
+++ b/tests/hf/test_upload.py
@@ -4,7 +4,6 @@
 
 import pytest
 import yaml
-
 from lanfactory.hf.upload import (
     DEFAULT_INCLUDE_PATTERNS,
     DEFAULT_REPO_ID,
diff --git a/tests/test_bayesflow_hssm_integration.py b/tests/test_bayesflow_hssm_integration.py
index 066ba2c..4405882 100644
--- a/tests/test_bayesflow_hssm_integration.py
+++ b/tests/test_bayesflow_hssm_integration.py
@@ -24,23 +24,21 @@
 os.environ["KERAS_BACKEND"] = "torch"
 os.environ.setdefault("KERAS_TORCH_DEVICE", "cpu")
 
-from pathlib import Path  # noqa: E402
-
-import pytest  # noqa: E402
+from pathlib import Path
+
+import bayesflow as bf
+import keras
+import numpy as np
+import pandas as pd
+import pytest
+from bayesflow.datasets import OfflineDataset
+from bayesflow.networks.inference.coupling.transforms import (
+    AffineTransform,
+)
+from lanfactory.onnx import transform_bayesflow_to_onnx
+from ssms.basic_simulators.simulator import simulator
 
 hssm = pytest.importorskip("hssm")
-
-import numpy as np  # noqa: E402
-import pandas as pd  # noqa: E402
-
-import bayesflow as bf  # noqa: E402
-import keras  # noqa: E402
-from bayesflow.datasets import OfflineDataset  # noqa: E402
-from bayesflow.networks.inference.coupling.transforms import AffineTransform  # noqa: E402
-from ssms.basic_simulators.simulator import simulator  # noqa: E402
-
-from lanfactory.onnx import transform_bayesflow_to_onnx  # noqa: E402
-
 _DDM_PARAM_NAMES = ["v", "a", "z", "t"]
 _DDM_PARAM_LOW = np.array([-2.0, 0.6, 0.3, 0.1], dtype=np.float32)
 _DDM_PARAM_HIGH = np.array([2.0, 1.8, 0.7, 0.5], dtype=np.float32)
diff --git a/tests/test_bayesflow_nle_export.py b/tests/test_bayesflow_nle_export.py
index ede588c..2f27a76 100644
--- a/tests/test_bayesflow_nle_export.py
+++ b/tests/test_bayesflow_nle_export.py
@@ -17,29 +17,31 @@
 os.environ["KERAS_BACKEND"] = "torch"
 os.environ.setdefault("KERAS_TORCH_DEVICE", "cpu")
 
-from pathlib import Path  # noqa: E402
+from pathlib import Path
+from typing import ClassVar
+
+import bayesflow as bf
+import jax
+import jax.numpy as jnp
+import keras
+import numpy as np
+import onnxruntime as ort
+import pytest
+import torch
+from bayesflow.datasets import OfflineDataset
+from bayesflow.networks.inference.coupling.transforms import (
+    AffineTransform,
+)
+from jaxonnxruntime import call_onnx, config
+from lanfactory.onnx import transform_bayesflow_to_onnx
 
-import jax  # noqa: E402
+import onnx
+from tests._onnx_utils import max_int64_abs
 
 # Same reason as the sbi test: ONNX shape/index tensors are int64; JAX's default
 # int32 silently truncates them inside jaxonnxruntime translation.
 jax.config.update("jax_enable_x64", True)
 
-import jax.numpy as jnp  # noqa: E402
-import numpy as np  # noqa: E402
-import onnx  # noqa: E402
-import onnxruntime as ort  # noqa: E402
-import pytest  # noqa: E402
-import torch  # noqa: E402
-from jaxonnxruntime import call_onnx, config  # noqa: E402
-
-import bayesflow as bf  # noqa: E402
-import keras  # noqa: E402
-from bayesflow.datasets import OfflineDataset  # noqa: E402
-from bayesflow.networks.inference.coupling.transforms import AffineTransform  # noqa: E402
-
-from lanfactory.onnx import transform_bayesflow_to_onnx  # noqa: E402
-from tests._onnx_utils import max_int64_abs  # noqa: E402
 
 # bayesflow under KERAS_BACKEND=torch globally disables autograd at import to
 # avoid excessive memory in long training loops. Restore the global default so
@@ -400,7 +402,9 @@ def log_prob(self, samples, conditions):
             return (samples.sum() + conditions.sum()).reshape(1)
 
     class _Standardizer:
-        standardize_layers = {"inference_conditions": _FakeStandardizeLayer(_THETA_DIM)}
+        standardize_layers: ClassVar[dict[str, _FakeStandardizeLayer]] = {
+            "inference_conditions": _FakeStandardizeLayer(_THETA_DIM)
+        }
 
     class _Approx:
         adapter = None
diff --git a/tests/test_bayesflow_nre_export.py b/tests/test_bayesflow_nre_export.py
index b5258ab..ad014b1 100644
--- a/tests/test_bayesflow_nre_export.py
+++ b/tests/test_bayesflow_nre_export.py
@@ -14,25 +14,25 @@
 os.environ["KERAS_BACKEND"] = "torch"
 os.environ.setdefault("KERAS_TORCH_DEVICE", "cpu")
 
-from pathlib import Path  # noqa: E402
-
-import jax  # noqa: E402
+from pathlib import Path
+from typing import ClassVar
+
+import bayesflow as bf
+import jax
+import jax.numpy as jnp
+import keras
+import numpy as np
+import onnxruntime as ort
+import pytest
+import torch
+from bayesflow.datasets import OfflineDataset
+from jaxonnxruntime import call_onnx, config
+from lanfactory.onnx import transform_bayesflow_to_onnx
+
+import onnx
 
 jax.config.update("jax_enable_x64", True)
 
-import jax.numpy as jnp  # noqa: E402
-import numpy as np  # noqa: E402
-import onnx  # noqa: E402
-import onnxruntime as ort  # noqa: E402
-import pytest  # noqa: E402
-import torch  # noqa: E402
-from jaxonnxruntime import call_onnx, config  # noqa: E402
-
-import bayesflow as bf  # noqa: E402
-import keras  # noqa: E402
-from bayesflow.datasets import OfflineDataset  # noqa: E402
-
-from lanfactory.onnx import transform_bayesflow_to_onnx  # noqa: E402
 
 # bayesflow under KERAS_BACKEND=torch globally disables autograd at import to
 # avoid excessive memory in long training loops. Restore the global default so
@@ -261,7 +261,9 @@ def __call__(self, hidden):
             return hidden.sum().reshape(1, 1)
 
     class _Standardizer:
-        standardize_layers = {"inference_conditions": _FakeStandardizeLayer(_X_DIM)}
+        standardize_layers: ClassVar[dict[str, _FakeStandardizeLayer]] = {
+            "inference_conditions": _FakeStandardizeLayer(_X_DIM)
+        }
 
     class _Approx:
         def __init__(self):
diff --git a/tests/test_cli_utils.py b/tests/test_cli_utils.py
index fe806ff..38644af 100644
--- a/tests/test_cli_utils.py
+++ b/tests/test_cli_utils.py
@@ -1,10 +1,10 @@
 """Tests for CLI utilities."""
 
-from unittest.mock import patch, mock_open
+import pickle
 
 import pytest
-
-from lanfactory.cli.utils import _make_train_network_configs, _get_train_network_config
+import yaml
+from lanfactory.cli.utils import _get_train_network_config, _make_train_network_configs
 
 
 def test_make_train_network_configs_with_dict_args():
@@ -44,25 +44,30 @@ def test_make_train_network_configs_without_save_name():
 
 
 def test_make_train_network_configs_with_save_name(tmp_path):
-    """Test _make_train_network_configs with save_name (file written)."""
+    """Test _make_train_network_configs creates and writes the pickle file to disk."""
     save_name = "test_config.pickle"
+    expected_file_path = tmp_path / save_name
+
+    result = _make_train_network_configs(
+        training_data_folder="/fake/data",
+        train_val_split=0.9,
+        save_folder=str(tmp_path),
+        network_arg_dict=None,
+        train_arg_dict=None,
+        save_name=save_name,
+    )
+
+    assert result["config_file_name"] == expected_file_path.name
 
-    with patch("builtins.open", mock_open()), patch("pickle.dump") as mock_dump:
-        result = _make_train_network_configs(
-            training_data_folder="/fake/data",
-            train_val_split=0.9,
-            save_folder=str(tmp_path),
-            network_arg_dict=None,
-            train_arg_dict=None,
-            save_name=save_name,
-        )
+    assert expected_file_path.is_file()
 
-        assert mock_dump.called
-        assert result["config_file_name"] == tmp_path / save_name
+    with open(expected_file_path, "rb") as f:
+        saved_data = pickle.load(f)
+        assert saved_data is not None
 
 
-def test_get_train_network_config_lan():
-    """Test _get_train_network_config with LAN network type."""
+def test_get_train_network_config_lan(tmp_path):
+    """Test _get_train_network_config with LAN network type using real file I/O."""
     yaml_content = {
         "NETWORK_TYPE": "lan",
         "LAYER_SIZES": [[100, 100, 1]],
@@ -83,20 +88,20 @@ def test_get_train_network_config_lan():
         "MODEL": "ddm",
     }
 
-    with (
-        patch("builtins.open", mock_open()),
-        patch("yaml.safe_load", return_value=yaml_content),
-    ):
-        result = _get_train_network_config(yaml_config_path="fake.yaml", net_index=0)
+    yaml_file = tmp_path / "test_config.yaml"
+    with open(yaml_file, "w") as f:
+        yaml.dump(yaml_content, f)
 
-        assert result["config_dict"]["network_config"]["train_output_type"] == "logprob"
-        assert result["config_dict"]["train_config"]["loss"] == "huber"
-        assert result["config_dict"]["train_config"]["features_key"] == "lan_data"
-        assert result["config_dict"]["train_config"]["label_key"] == "lan_labels"
-        assert result["extra_fields"]["model"] == "ddm"
+    result = _get_train_network_config(yaml_config_path=yaml_file, net_index=0)
 
+    assert result["config_dict"]["network_config"]["train_output_type"] == "logprob"
+    assert result["config_dict"]["train_config"]["loss"] == "huber"
+    assert result["config_dict"]["train_config"]["features_key"] == "lan_data"
+    assert result["config_dict"]["train_config"]["label_key"] == "lan_labels"
+    assert result["extra_fields"]["model"] == "ddm"
 
-def test_get_train_network_config_cpn():
+
+def test_get_train_network_config_cpn(tmp_path):
     """Test _get_train_network_config with CPN network type."""
     yaml_content = {
         "NETWORK_TYPE": "cpn",
@@ -118,19 +123,18 @@ def test_get_train_network_config_cpn():
         "MODEL": "ddm",
     }
 
-    with (
-        patch("builtins.open", mock_open()),
-        patch("yaml.safe_load", return_value=yaml_content),
-    ):
-        result = _get_train_network_config(yaml_config_path="fake.yaml", net_index=0)
+    config_file = tmp_path / "test_config.yaml"
+    config_file.write_text(yaml.dump(yaml_content))
+
+    result = _get_train_network_config(yaml_config_path=config_file, net_index=0)
 
-        assert result["config_dict"]["network_config"]["train_output_type"] == "logits"
-        assert result["config_dict"]["train_config"]["loss"] == "bcelogit"
-        assert result["config_dict"]["train_config"]["features_key"] == "cpn_data"
-        assert result["config_dict"]["train_config"]["label_key"] == "cpn_labels"
+    assert result["config_dict"]["network_config"]["train_output_type"] == "logits"
+    assert result["config_dict"]["train_config"]["loss"] == "bcelogit"
+    assert result["config_dict"]["train_config"]["features_key"] == "cpn_data"
+    assert result["config_dict"]["train_config"]["label_key"] == "cpn_labels"
 
 
-def test_get_train_network_config_opn():
+def test_get_train_network_config_opn(tmp_path):
     """Test _get_train_network_config with OPN network type."""
     yaml_content = {
         "NETWORK_TYPE": "opn",
@@ -152,16 +156,15 @@ def test_get_train_network_config_opn():
         "MODEL": "ddm",
     }
 
-    with (
-        patch("builtins.open", mock_open()),
-        patch("yaml.safe_load", return_value=yaml_content),
-    ):
-        result = _get_train_network_config(yaml_config_path="fake.yaml", net_index=0)
+    config_file = tmp_path / "test_config_opn.yaml"
+    config_file.write_text(yaml.dump(yaml_content))
 
-        assert result["config_dict"]["network_config"]["train_output_type"] == "logits"
-        assert result["config_dict"]["train_config"]["loss"] == "bcelogit"
-        assert result["config_dict"]["train_config"]["features_key"] == "opn_data"
-        assert result["config_dict"]["train_config"]["label_key"] == "opn_labels"
+    result = _get_train_network_config(yaml_config_path=config_file, net_index=0)
+
+    assert result["config_dict"]["network_config"]["train_output_type"] == "logits"
+    assert result["config_dict"]["train_config"]["loss"] == "bcelogit"
+    assert result["config_dict"]["train_config"]["features_key"] == "opn_data"
+    assert result["config_dict"]["train_config"]["label_key"] == "opn_labels"
 
 
 def test_get_train_network_config_no_path():
@@ -170,7 +173,7 @@ def test_get_train_network_config_no_path():
         _get_train_network_config(yaml_config_path=None)
 
 
-def test_get_train_network_config_with_net_index():
+def test_get_train_network_config_with_net_index(tmp_path):
     """Test _get_train_network_config with different net_index."""
     yaml_content = {
         "NETWORK_TYPE": "lan",
@@ -192,17 +195,16 @@ def test_get_train_network_config_with_net_index():
         "MODEL": "ddm",
     }
 
-    with (
-        patch("builtins.open", mock_open()),
-        patch("yaml.safe_load", return_value=yaml_content),
-    ):
-        result = _get_train_network_config(yaml_config_path="fake.yaml", net_index=1)
-
-        # layer_sizes comes directly from YAML (not modified)
-        assert result["config_dict"]["network_config"]["layer_sizes"] == [120, 120, 1]
-        # activations has output layer activation appended
-        assert result["config_dict"]["network_config"]["activations"] == [
-            "relu",
-            "relu",
-            "linear",
-        ]
+    config_file = tmp_path / "test_config_net_index.yaml"
+    config_file.write_text(yaml.dump(yaml_content))
+
+    result = _get_train_network_config(yaml_config_path=config_file, net_index=1)
+
+    # layer_sizes comes directly from YAML (not modified)
+    assert result["config_dict"]["network_config"]["layer_sizes"] == [120, 120, 1]
+    # activations has output layer activation appended
+    assert result["config_dict"]["network_config"]["activations"] == [
+        "relu",
+        "relu",
+        "linear",
+    ]
diff --git a/tests/test_end_to_end_jax.py b/tests/test_end_to_end_jax.py
index 9a5a65d..dd84904 100644
--- a/tests/test_end_to_end_jax.py
+++ b/tests/test_end_to_end_jax.py
@@ -1,19 +1,19 @@
-import pytest
-import ssms
-import lanfactory
+import logging
 import os
-import numpy as np
 from copy import deepcopy
+
 import jax.numpy as jnp
+import lanfactory
+import numpy as np
+import pytest
+import ssms
 import torch
+
 from .constants import (
     TEST_GENERATOR_CONSTANTS,
     TEST_MODEL_FOLDER_CONSTANTS_JAX,
 )
 
-# import logger
-import logging
-
 logger = logging.getLogger(__name__)
 
 LEN_FORWARD_PASS_DUMMY = 2000
diff --git a/tests/test_end_to_end_torch.py b/tests/test_end_to_end_torch.py
index 22a82ae..8a5c5c9 100644
--- a/tests/test_end_to_end_torch.py
+++ b/tests/test_end_to_end_torch.py
@@ -1,17 +1,18 @@
-import pytest
-import ssms
-import lanfactory
+import logging
 import os
-import numpy as np
 from copy import deepcopy
+
+import lanfactory
+import numpy as np
+import pytest
+import ssms
 import torch
+
 from .constants import (
     TEST_GENERATOR_CONSTANTS,
     TEST_MODEL_FOLDER_CONSTANTS_TORCH,
 )
 
-import logging
-
 logger = logging.getLogger(__name__)
 
 LEN_FORWARD_PASS_DUMMY = 2000
diff --git a/tests/test_jax_mlp.py b/tests/test_jax_mlp.py
index 895a1e6..163c2bb 100644
--- a/tests/test_jax_mlp.py
+++ b/tests/test_jax_mlp.py
@@ -2,11 +2,10 @@
 
 import pickle
 
-import pytest
 import jax
 import jax.numpy as jnp
-
-from lanfactory.trainers.jax_mlp import JaxMLPFactory, JaxMLP
+import pytest
+from lanfactory.trainers.jax_mlp import JaxMLP, JaxMLPFactory
 
 
 def test_mlp_jax_factory_with_dict():
@@ -44,9 +43,9 @@ def test_mlp_jax_factory_with_string_path(tmp_path):
     assert model.layer_sizes == [100, 100, 1]
 
 
-def test_mlp_jax_factory_raises_value_error():
+def test_mlp_jax_factory_raises_type_error():
     """Test JaxMLPFactory raises ValueError for invalid network_config type."""
-    with pytest.raises(ValueError, match="network_config argument is not passed"):
+    with pytest.raises(TypeError, match="network_config argument is not passed"):
         JaxMLPFactory(network_config=123, train=True)  # Invalid type
 
 
@@ -128,7 +127,7 @@ def test_mlp_jax_forward_with_non_linear_output_activation():
     x = jax.random.normal(key, (5, 5))
 
     # Initialize model state
-    key1, key2 = jax.random.split(key)
+    key1, _ = jax.random.split(key)
     state = model.init(key1, x)
 
     # Forward pass
@@ -152,7 +151,7 @@ def test_mlp_jax_inference_mode_with_logits():
     x = jax.random.normal(key, (5, 5))
 
     # Initialize model state
-    key1, key2 = jax.random.split(key)
+    key1, _ = jax.random.split(key)
     state = model.init(key1, x)
 
     # Forward pass in inference mode
@@ -163,8 +162,8 @@ def test_mlp_jax_inference_mode_with_logits():
 
 def test_mlp_jax_load_state_from_file_error():
     """Test JaxMLP load_state_from_file raises error when file_path is None."""
-    from lanfactory.trainers.jax_mlp import JaxMLPFactory
     import pytest
+    from lanfactory.trainers.jax_mlp import JaxMLPFactory
 
     network_config = {
         "layer_sizes": [10, 10, 1],
@@ -180,8 +179,8 @@ def test_mlp_jax_load_state_from_file_error():
 
 def test_mlp_jax_load_state_from_file_without_input_dim(tmp_path):
     """Test JaxMLP load_state_from_file without providing input_dim."""
-    from lanfactory.trainers.jax_mlp import JaxMLPFactory
     import flax.serialization
+    from lanfactory.trainers.jax_mlp import JaxMLPFactory
 
     network_config = {
         "layer_sizes": [10, 10, 1],
@@ -242,8 +241,8 @@ def test_mlp_jax_make_forward_partial_with_dict_state(tmp_path):
 
 def test_mlp_jax_make_forward_partial_without_jit(tmp_path):
     """Test JaxMLP make_forward_partial without JIT compilation."""
-    from lanfactory.trainers.jax_mlp import JaxMLPFactory
     import flax.serialization
+    from lanfactory.trainers.jax_mlp import JaxMLPFactory
 
     network_config = {
         "layer_sizes": [10, 10, 1],
@@ -276,8 +275,8 @@ def test_mlp_jax_make_forward_partial_without_jit(tmp_path):
 
 def test_mlp_jax_make_forward_partial_invalid_state_type():
     """Test JaxMLP make_forward_partial raises error with invalid state type."""
-    from lanfactory.trainers.jax_mlp import JaxMLPFactory
     import pytest
+    from lanfactory.trainers.jax_mlp import JaxMLPFactory
 
     network_config = {
         "layer_sizes": [10, 10, 1],
@@ -289,7 +288,7 @@ def test_mlp_jax_make_forward_partial_invalid_state_type():
 
     # Test with invalid state type (list instead of dict or string)
     with pytest.raises(
-        ValueError, match="state argument has to be a dictionary or a string"
+        TypeError, match="state argument has to be a dictionary or a string"
     ):
         model.make_forward_partial(
             seed=42, input_dim=5, state=[1, 2, 3], add_jitted=True
diff --git a/tests/test_mlflow_integration.py b/tests/test_mlflow_integration.py
index c2f348a..7d3ed07 100644
--- a/tests/test_mlflow_integration.py
+++ b/tests/test_mlflow_integration.py
@@ -4,8 +4,8 @@
 import shutil
 from pathlib import Path
 
-import pytest
 import numpy as np
+import pytest
 
 try:
     import mlflow
@@ -51,7 +51,7 @@ def cleanup_mlflow():
     try:
         mlflow.set_tracking_uri(original_uri)
     except Exception:
-        pass
+        mlflow.set_tracking_uri("file:./mlruns")
 
 
 @pytest.fixture
@@ -306,9 +306,9 @@ def test_jax_trainer_mlflow_logging(
         tmp_path,
     ):
         """Test that JAX trainer logs to MLflow correctly."""
-        from torch.utils.data import DataLoader
         from lanfactory.trainers.jax_mlp import JaxMLPFactory, ModelTrainerJaxMLP
         from lanfactory.trainers.torch_mlp import DatasetTorch
+        from torch.utils.data import DataLoader
 
         tracking_uri = test_mlflow_dir["tracking_uri"]
         artifact_location = test_mlflow_dir["artifact_location"]
@@ -400,12 +400,12 @@ def test_pytorch_trainer_mlflow_logging(
         tmp_path,
     ):
         """Test that PyTorch trainer logs to MLflow correctly."""
-        from torch.utils.data import DataLoader
         from lanfactory.trainers.torch_mlp import (
-            TorchMLP,
-            ModelTrainerTorchMLP,
             DatasetTorch,
+            ModelTrainerTorchMLP,
+            TorchMLP,
         )
+        from torch.utils.data import DataLoader
 
         tracking_uri = test_mlflow_dir["tracking_uri"]
         artifact_location = test_mlflow_dir["artifact_location"]
@@ -639,12 +639,12 @@ def test_trainer_without_mlflow(
         tmp_path,
     ):
         """Test that trainers work correctly when MLflow is disabled."""
-        from torch.utils.data import DataLoader
         from lanfactory.trainers.torch_mlp import (
-            TorchMLP,
-            ModelTrainerTorchMLP,
             DatasetTorch,
+            ModelTrainerTorchMLP,
+            TorchMLP,
         )
+        from torch.utils.data import DataLoader
 
         # Generate minimal training data
         gen_configs = dummy_generator_config_simple_two_choices()
diff --git a/tests/test_network_inspectors_api.py b/tests/test_network_inspectors_api.py
index 606eb95..fcd8c88 100644
--- a/tests/test_network_inspectors_api.py
+++ b/tests/test_network_inspectors_api.py
@@ -5,7 +5,6 @@
 import numpy as np
 import pandas as pd
 import pytest
-
 from lanfactory.network_inspectors import api
 
 
diff --git a/tests/test_network_inspectors_plotting.py b/tests/test_network_inspectors_plotting.py
index 4498ad0..5d35f98 100644
--- a/tests/test_network_inspectors_plotting.py
+++ b/tests/test_network_inspectors_plotting.py
@@ -5,7 +5,6 @@
 import numpy as np
 import pandas as pd
 import plotly.graph_objects as go
-
 from lanfactory.network_inspectors.config import ModelSpec, PlotConfig
 from lanfactory.network_inspectors.plotting import plot_manifold
 
diff --git a/tests/test_sbi_embeddings.py b/tests/test_sbi_embeddings.py
index b39a961..2c390fd 100644
--- a/tests/test_sbi_embeddings.py
+++ b/tests/test_sbi_embeddings.py
@@ -13,23 +13,22 @@
 from pathlib import Path
 
 import jax
+import numpy as np
+import onnxruntime as ort
+import pytest
+import torch
+from jaxonnxruntime import call_onnx, config
+from lanfactory.onnx import transform_sbi_to_onnx
+from sbi.inference import NRE_A
+from sbi.neural_nets import classifier_nn
+from sbi.neural_nets.embedding_nets import CNNEmbedding, FCEmbedding
+from sbi.utils import BoxUniform
+from torch import nn
+
+import onnx
 
 jax.config.update("jax_enable_x64", True)
 
-import numpy as np  # noqa: E402
-import onnx  # noqa: E402
-import onnxruntime as ort  # noqa: E402
-import pytest  # noqa: E402
-import torch  # noqa: E402
-from jaxonnxruntime import call_onnx, config  # noqa: E402
-from sbi.inference import NRE_A  # noqa: E402
-from sbi.neural_nets import classifier_nn  # noqa: E402
-from sbi.neural_nets.embedding_nets import CNNEmbedding, FCEmbedding  # noqa: E402
-from sbi.utils import BoxUniform  # noqa: E402
-from torch import nn  # noqa: E402
-
-from lanfactory.onnx import transform_sbi_to_onnx  # noqa: E402
-
 config.update("jaxort_only_allow_initializers_as_static_args", False)
 
 # sbi's build_mlp_classifier defaults to nn.LayerNorm between hidden layers, but
diff --git a/tests/test_sbi_hssm_integration.py b/tests/test_sbi_hssm_integration.py
index fcf34ba..97b41ab 100644
--- a/tests/test_sbi_hssm_integration.py
+++ b/tests/test_sbi_hssm_integration.py
@@ -26,14 +26,13 @@
 # Skip cleanly when HSSM is not in the environment.
 hssm = pytest.importorskip("hssm")
 
-import numpy as np  # noqa: E402
-import pandas as pd  # noqa: E402
-import torch  # noqa: E402
-from sbi.inference import NLE_A  # noqa: E402
-from sbi.utils import BoxUniform  # noqa: E402
-from ssms.basic_simulators.simulator import simulator  # noqa: E402
-
-from lanfactory.onnx import transform_sbi_to_onnx  # noqa: E402
+import numpy as np
+import pandas as pd
+import torch
+from lanfactory.onnx import transform_sbi_to_onnx
+from sbi.inference import NLE_A
+from sbi.utils import BoxUniform
+from ssms.basic_simulators.simulator import simulator
 
 # DDM parameter order matches sbi simulator inputs and HSSM defaults.
 _DDM_PARAM_NAMES = ["v", "a", "z", "t"]
diff --git a/tests/test_sbi_nle_export.py b/tests/test_sbi_nle_export.py
index d2a41fc..52bb772 100644
--- a/tests/test_sbi_nle_export.py
+++ b/tests/test_sbi_nle_export.py
@@ -17,18 +17,18 @@
 # wrong numerical values (~0.5 drift from the torch reference on MAF log_prob).
 jax.config.update("jax_enable_x64", True)
 
-import jax.numpy as jnp  # noqa: E402
-import numpy as np  # noqa: E402
-import onnx  # noqa: E402
-import onnxruntime as ort  # noqa: E402
-import pytest  # noqa: E402
-import torch  # noqa: E402
-from jaxonnxruntime import call_onnx, config  # noqa: E402
-from sbi.inference import NLE_A  # noqa: E402
-from sbi.utils import BoxUniform  # noqa: E402
-
-from lanfactory.onnx import transform_sbi_to_onnx  # noqa: E402
-from tests._onnx_utils import max_int64_abs  # noqa: E402
+import jax.numpy as jnp
+import numpy as np
+import onnxruntime as ort
+import pytest
+import torch
+from jaxonnxruntime import call_onnx, config
+from lanfactory.onnx import transform_sbi_to_onnx
+from sbi.inference import NLE_A
+from sbi.utils import BoxUniform
+
+import onnx
+from tests._onnx_utils import max_int64_abs
 
 # Same friction as C2's MAF spike — torch.onnx.export emits Reshape shapes as
 # Constant nodes. HSSM's onnx2jax patch (commit 2e76516) sets this globally for
@@ -191,7 +191,7 @@ def test_nle_log_prob_ordering_matches_analytical_gaussian(
 def test_transform_rejects_unsupported_score_estimator(tmp_path: Path) -> None:
     """Estimators in the unsupported set should fail loudly."""
 
-    class ScoreEstimator(torch.nn.Module):  # noqa: D401 - name is the signal
+    class ScoreEstimator(torch.nn.Module):
         pass
 
     with pytest.raises(ValueError, match="does not support"):
diff --git a/tests/test_sbi_nre_export.py b/tests/test_sbi_nre_export.py
index 20a7433..7ced473 100644
--- a/tests/test_sbi_nre_export.py
+++ b/tests/test_sbi_nre_export.py
@@ -14,17 +14,17 @@
 # x64 required before any JAX import — see test_sbi_nle_export.py for details.
 jax.config.update("jax_enable_x64", True)
 
-import jax.numpy as jnp  # noqa: E402
-import numpy as np  # noqa: E402
-import onnx  # noqa: E402
-import onnxruntime as ort  # noqa: E402
-import pytest  # noqa: E402
-import torch  # noqa: E402
-from jaxonnxruntime import call_onnx, config  # noqa: E402
-from sbi.inference import NRE_A  # noqa: E402
-from sbi.utils import BoxUniform  # noqa: E402
-
-from lanfactory.onnx import transform_sbi_to_onnx  # noqa: E402
+import jax.numpy as jnp
+import numpy as np
+import onnxruntime as ort
+import pytest
+import torch
+from jaxonnxruntime import call_onnx, config
+from lanfactory.onnx import transform_sbi_to_onnx
+from sbi.inference import NRE_A
+from sbi.utils import BoxUniform
+
+import onnx
 
 config.update("jaxort_only_allow_initializers_as_static_args", False)
 
diff --git a/tests/test_sbi_spike_maf_roundtrip.py b/tests/test_sbi_spike_maf_roundtrip.py
index f0eb174..dec708e 100644
--- a/tests/test_sbi_spike_maf_roundtrip.py
+++ b/tests/test_sbi_spike_maf_roundtrip.py
@@ -11,7 +11,6 @@
 
 import jax
 import numpy as np
-import onnx
 import onnxruntime as ort
 import pytest
 import torch
@@ -25,6 +24,8 @@
 )
 from torch import nn
 
+import onnx
+
 # Friction discovered in C2: nflows' MAF exports a Reshape whose shape argument
 # is a Constant node (not a model initializer). jaxonnxruntime's default strict
 # mode rejects this. The flag below tells jaxonnxruntime to treat Constant nodes
diff --git a/tests/test_sbi_spike_mlp_roundtrip.py b/tests/test_sbi_spike_mlp_roundtrip.py
index 100beb0..652de04 100644
--- a/tests/test_sbi_spike_mlp_roundtrip.py
+++ b/tests/test_sbi_spike_mlp_roundtrip.py
@@ -10,13 +10,14 @@
 
 import jax
 import numpy as np
-import onnx
 import onnxruntime as ort
 import pytest
 import torch
 from jaxonnxruntime import call_onnx
 from torch import nn
 
+import onnx
+
 
 @pytest.mark.flaky(reruns=2)
 def test_mlp_three_way_agreement(tmp_path: Path) -> None:
diff --git a/tests/test_torch_mlp.py b/tests/test_torch_mlp.py
index da38fe9..8076907 100644
--- a/tests/test_torch_mlp.py
+++ b/tests/test_torch_mlp.py
@@ -1,15 +1,15 @@
 """Tests for the DatasetTorch class and related components."""
 
 import pickle
+from unittest.mock import MagicMock
+
 import numpy as np
 import pytest
 import torch
-from unittest.mock import MagicMock
-
 from lanfactory.trainers.torch_mlp import (
     DatasetTorch,
-    ModelTrainerTorchMLP,
     LoadTorchMLPInfer,
+    ModelTrainerTorchMLP,
     TorchMLP,
 )
 
@@ -110,7 +110,7 @@ def test_dataset_torch_getitem_loads_file_on_first_access(
     # First access should load file and populate tmp_data
     assert dataset.tmp_data == {}  # Initially empty
 
-    X, y = dataset[0]
+    _, _ = dataset[0]
 
     # Now tmp_data should be populated
     assert dataset.tmp_data != {}
@@ -132,9 +132,9 @@ def test_dataset_torch_getitem_multiple_batches_same_file(
     )
 
     # Access multiple batches from the same file
-    X0, y0 = dataset[0]
-    X1, y1 = dataset[1]
-    X2, y2 = dataset[2]
+    X0, _y0 = dataset[0]
+    X1, _y1 = dataset[1]
+    X2, _y2 = dataset[2]
 
     # All should have correct shape
     assert X0.shape == (200, 6)
@@ -163,22 +163,22 @@ def test_dataset_torch_getitem_crosses_file_boundary(
     # So indices 0-3 are from file 0, 4-7 from file 1, 8-11 from file 2
 
     # Access batch from first file
-    X0, y0 = dataset[0]
+    _X0, _y0 = dataset[0]
     first_file_data = dataset.tmp_data["lan_data"].copy()
 
     # Access batch from same file
-    X3, y3 = dataset[3]
+    _X3, _y3 = dataset[3]
     assert np.array_equal(dataset.tmp_data["lan_data"], first_file_data)
 
     # Access first batch from second file - should trigger file load
-    X4, y4 = dataset[4]
+    _X4, _y4 = dataset[4]
     second_file_data = dataset.tmp_data["lan_data"]
 
     # Data should be different (new file loaded)
     assert not np.array_equal(second_file_data, first_file_data)
 
     # Access first batch from third file
-    X8, y8 = dataset[8]
+    _X8, _y8 = dataset[8]
     third_file_data = dataset.tmp_data["lan_data"]
 
     # Should be different from second file
@@ -200,7 +200,7 @@ def test_dataset_torch_getitem_with_label_bounds(
         label_key="lan_labels",
     )
 
-    X, y = dataset[0]
+    _, y = dataset[0]
 
     # Labels should be clipped to bounds
     assert np.all(y >= -10.0)
@@ -271,7 +271,7 @@ def test_dataset_torch_empty_tmp_data_triggers_load(
     )
 
     # Access batch 2 (not 0), but tmp_data is empty
-    X, y = dataset[2]
+    X, _y = dataset[2]
 
     # Should work - file was loaded due to empty tmp_data
     assert X.shape == (200, 6)
@@ -342,7 +342,7 @@ def test_dataset_torch_label_bounds(tmp_path):
         label_key="lan_labels",
     )
 
-    X, y = dataset[0]
+    _, y = dataset[0]
 
     # Labels should be clipped to bounds
     assert np.all(y >= -10.0)
@@ -367,7 +367,7 @@ def test_dataset_torch_3d_labels_raises_error(tmp_path):
     )
 
     with pytest.raises(ValueError, match="Label data has unexpected shape"):
-        X, y = dataset[0]
+        _X, _y = dataset[0]
 
 
 def test_dataset_torch_batch_size_not_divisible_raises_error(tmp_path):
@@ -987,12 +987,12 @@ def test_torch_mlp_with_non_linear_output_activation():
 
 def test_model_trainer_torch_mlp_with_none_train_config(create_mock_data_files):
     """Test ModelTrainerTorchMLP raises error when train_config is None."""
+    import pytest
     from lanfactory.trainers.torch_mlp import (
         DatasetTorch,
         ModelTrainerTorchMLP,
         TorchMLP,
     )
-    import pytest
 
     file_list = create_mock_data_files(n_files=1)
 
@@ -1113,7 +1113,7 @@ def test_make_train_valid_dataloaders_raises_no_valid_files(create_mock_data_fil
 
 def test_torch_mlp_factory_with_dict():
     """Test TorchMLPFactory with dict config."""
-    from lanfactory.trainers.torch_mlp import TorchMLPFactory, TorchMLP
+    from lanfactory.trainers.torch_mlp import TorchMLP, TorchMLPFactory
 
     network_config = {
         "layer_sizes": [10, 10, 1],
@@ -1129,7 +1129,7 @@ def test_torch_mlp_factory_with_dict():
 
 def test_torch_mlp_factory_with_pickle_path(tmp_path):
     """Test TorchMLPFactory with path to pickled config."""
-    from lanfactory.trainers.torch_mlp import TorchMLPFactory, TorchMLP
+    from lanfactory.trainers.torch_mlp import TorchMLP, TorchMLPFactory
 
     network_config = {
         "layer_sizes": [10, 10, 1],
diff --git a/tests/test_transform_onnx.py b/tests/test_transform_onnx.py
index ec663e2..451c9fd 100644
--- a/tests/test_transform_onnx.py
+++ b/tests/test_transform_onnx.py
@@ -5,8 +5,7 @@
 
 import pytest
 import torch
-
-from lanfactory.onnx.transform_onnx import transform_to_onnx, main
+from lanfactory.onnx.transform_onnx import main, transform_to_onnx
 
 
 @pytest.fixture
@@ -165,14 +164,16 @@ def test_transform_to_onnx_creates_correct_input_tensor(mock_network_config):
 def test_transform_to_onnx_missing_config_file():
     """Test that transform_to_onnx raises error for missing config file."""
     # Mock open to raise FileNotFoundError
-    with patch("builtins.open", side_effect=FileNotFoundError("File not found")):
-        with pytest.raises(FileNotFoundError):
-            transform_to_onnx(
-                network_config_file="/nonexistent/config.pickle",
-                state_dict_file="/fake/state.pt",
-                input_shape=6,
-                output_onnx_file="/fake/output.onnx",
-            )
+    with (
+        patch("builtins.open", side_effect=FileNotFoundError("File not found")),
+        pytest.raises(FileNotFoundError),
+    ):
+        transform_to_onnx(
+            network_config_file="/nonexistent/config.pickle",
+            state_dict_file="/fake/state.pt",
+            input_shape=6,
+            output_onnx_file="/fake/output.onnx",
+        )
 
 
 def test_transform_to_onnx_missing_state_dict_file(mock_network_config):
@@ -210,14 +211,14 @@ def test_transform_to_onnx_invalid_pickle_file():
             "lanfactory.onnx.transform_onnx.pickle.load",
             side_effect=pickle.UnpicklingError("Invalid pickle"),
         ),
+        pytest.raises(pickle.UnpicklingError),
     ):
-        with pytest.raises(pickle.UnpicklingError):
-            transform_to_onnx(
-                network_config_file=config_file,
-                state_dict_file="/fake/state.pt",
-                input_shape=6,
-                output_onnx_file="/fake/output.onnx",
-            )
+        transform_to_onnx(
+            network_config_file=config_file,
+            state_dict_file="/fake/state.pt",
+            input_shape=6,
+            output_onnx_file="/fake/output.onnx",
+        )
 
 
 def test_main_calls_transform_to_onnx():
diff --git a/tests/utils.py b/tests/utils.py
index 678a944..d177b0d 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,7 +1,6 @@
-import shutil
-import pathlib
-
 import logging
+import pathlib
+import shutil
 
 logger = logging.getLogger(__name__)
 
@@ -52,7 +51,7 @@ def clean_out_folder(folder: str | pathlib.Path | None = None, dry_run=True) ->
         except PermissionError:
             logger.error(f"Permission denied when trying to remove folder '{folder}'.")
         except Exception as e:
-            logger.error(f"Error removing folder '{folder}': {str(e)}")
+            logger.error(f"Error removing folder '{folder}': {e!s}")
     else:
         logger.error(f"Folder '{folder}' does not exist.")
 
@@ -61,7 +60,7 @@ def print_tree(
     path: pathlib.Path | str,
     prefix: str = "",
     logger: logging.Logger | None = logger,
-    out_str_list: list[str] = [],
+    out_str_list: list[str] | None = None,
 ) -> list[str]:
     """Print a directory tree structure starting from the given path.
 
@@ -86,6 +85,8 @@ def print_tree(
         print_tree(Path("./my_directory"))
     """
     path = pathlib.Path(path)
+    if out_str_list is None:
+        out_str_list = []
     contents = sorted(path.iterdir(), key=lambda p: (not p.is_dir(), p.name.lower()))
     for index, item in enumerate(contents):
         connector = "└── " if index == len(contents) - 1 else "├── "