Repository: aleximmer/Laplace (language: python)

Issue:
Add linting and formatting via `ruff`
Use [ruff](https://github.com/astral-sh/ruff) for linting and formatting.
Fix all linting issues and format all files with `ruff`
Follow-up for #176. We should wait until #144 is merged also.


Discussion / hints:
I've been personally using `ruff`, for both linting and formatting, for a while. Its default setting is satisfactory. So let's go ahead with this.

I guess we still need to decide if we want single quote or double quote for string literals. But this is easy for `ruff`---currently in the `mc-subset2` branch I have the following config, to conform to the current convention:
https://github.com/aleximmer/Laplace/blob/bd5501cd96e7f3bc0d189b4a9632429f6eef31ec/pyproject.toml#L8-L9


Fix patch:
diff --git a/.github/workflows/format-ruff.yml b/.github/workflows/format-ruff.yml
new file mode 100644
index 0000000..50190fc
--- /dev/null
+++ b/.github/workflows/format-ruff.yml
@@ -0,0 +1,12 @@
+name: ruff-format
+
+on: pull_request
+
+jobs:
+  ruff-format:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: chartboost/ruff-action@v1
+        with:
+          args: "format --check"
diff --git a/.github/workflows/lint-ruff.yml b/.github/workflows/lint-ruff.yml
new file mode 100644
index 0000000..be99ccf
--- /dev/null
+++ b/.github/workflows/lint-ruff.yml
@@ -0,0 +1,12 @@
+name: ruff-lint
+
+on: pull_request
+
+jobs:
+  ruff-lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: chartboost/ruff-action@v1
+        with:
+          args: 'check'
diff --git a/README.md b/README.md
index 343543e..8ad27ce 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,8 @@ The laplace package facilitates the application of Laplace approximations for en
 The package enables posterior approximations, marginal-likelihood estimation, and various posterior predictive computations.
 The library documentation is available at [https://aleximmer.github.io/Laplace](https://aleximmer.github.io/Laplace).
 
-There is also a corresponding paper, [*Laplace Redux — Effortless Bayesian Deep Learning*](https://arxiv.org/abs/2106.14806), which introduces the library, provides an introduction to the Laplace approximation, reviews its use in deep learning, and empirically demonstrates its versatility and competitiveness. Please consider referring to the paper when using our library:
+There is also a corresponding paper, [_Laplace Redux — Effortless Bayesian Deep Learning_](https://arxiv.org/abs/2106.14806), which introduces the library, provides an introduction to the Laplace approximation, reviews its use in deep learning, and empirically demonstrates its versatility and competitiveness. Please consider referring to the paper when using our library:
+
 ```bibtex
 @inproceedings{laplace2021,
   title={Laplace Redux--Effortless {B}ayesian Deep Learning},
@@ -18,19 +19,22 @@ There is also a corresponding paper, [*Laplace Redux — Effortless Bayesian Dee
   year={2021}
 }
 ```
+
 The [code](https://github.com/runame/laplace-redux) to reproduce the experiments in the paper is also publicly available; it provides examples of how to use our library for predictive uncertainty quantification, model selection, and continual learning.
 
 ## Setup
 
 For full compatibility, install this package in a fresh virtual env.
-We assume Python >= 3.8 (last tested on Python 3.10.12).
+We assume Python >= 3.9 since lower versions are [(soon to be) deprecated](https://devguide.python.org/versions/).
 PyTorch version 2.0 and up is also required for full compatibility.
 To install laplace with `pip`, run the following:
+
 ```bash
 pip install laplace-torch
 ```
 
 For development purposes, clone the repository and then install:
+
 ```bash
 # or after cloning the repository for development
 pip install -e .
@@ -41,14 +45,14 @@ pytest tests/
 
 ## Example usage
 
-### *Post-hoc* prior precision tuning of diagonal LA
+### _Post-hoc_ prior precision tuning of diagonal LA
 
 In the following example, a pre-trained model is loaded,
 then the Laplace approximation is fit to the training data
 (using a diagonal Hessian approximation over all parameters),
-and the prior precision is optimized with cross-validation `'gridsearch'`.
+and the prior precision is optimized with cross-validation `"gridsearch"`.
 After that, the resulting LA is used for prediction with
-the `'probit'` predictive for classification.
+the `"probit"` predictive for classification.
 
 ```python
 from laplace import Laplace
@@ -57,14 +61,14 @@ from laplace import Laplace
 model = load_map_model()
 
 # User-specified LA flavor
-la = Laplace(model, 'classification',
-             subset_of_weights='all',
-             hessian_structure='diag')
+la = Laplace(model, "classification",
+             subset_of_weights="all",
+             hessian_structure="diag")
 la.fit(train_loader)
-la.optimize_prior_precision(method='gridsearch', val_loader=val_loader)
+la.optimize_prior_precision(method="gridsearch", val_loader=val_loader)
 
 # User-specified predictive approx.
-pred = la(x, link_approx='probit')
+pred = la(x, link_approx="probit")
 ```
 
 ### Differentiating the log marginal likelihood w.r.t. hyperparameters
@@ -81,7 +85,7 @@ from laplace import Laplace
 model = load_model()
 
 # Default to recommended last-layer KFAC LA:
-la = Laplace(model, likelihood='regression')
+la = Laplace(model, likelihood="regression")
 la.fit(train_loader)
 
 # ML w.r.t. prior precision and observation noise
@@ -98,15 +102,15 @@ See `examples/huggingface_examples.py` for a runnable example.
 class MyGPT2(nn.Module):
     def __init__(self, tokenizer: PreTrainedTokenizer) -> None:
         super().__init__()
-        config = GPT2Config.from_pretrained('gpt2')
+        config = GPT2Config.from_pretrained("gpt2")
         config.pad_token_id = tokenizer.pad_token_id
         config.num_labels = 2
         self.hf_model = GPT2ForSequenceClassification.from_pretrained(
-            'gpt2', config=config
+            "gpt2", config=config
         )
 
     def forward(self, data: MutableMapping) -> torch.Tensor:
-        '''
+        """
         Custom forward function. Handles things like moving the
         input tensor to the correct device inside.
 
@@ -117,10 +121,10 @@ class MyGPT2(nn.Module):
 
         Returns:
             logits: An `(batch_size, n_classes)`-sized tensor of logits.
-        '''
+        """
         device = next(self.parameters()).device
-        input_ids = data['input_ids'].to(device)
-        attn_mask = data['attention_mask'].to(device)
+        input_ids = data["input_ids"].to(device)
+        attn_mask = data["attention_mask"].to(device)
         output_dict = self.hf_model(input_ids=input_ids, attention_mask=attn_mask)
         return output_dict.logits
 
@@ -137,10 +141,10 @@ for p in model.hf_model.score.parameters():
 
 la = Laplace(
     model,
-    likelihood='classification',
+    likelihood="classification",
     # Will only hit the last-layer since it's the only one that is grad-enabled
-    subset_of_weights='all',
-    hessian_structure='diag',
+    subset_of_weights="all",
+    hessian_structure="diag",
 )
 la.fit(dataloader)
 la.optimize_prior_precision()
@@ -154,9 +158,9 @@ def get_lora_model():
     config = LoraConfig(
         r=4,
         lora_alpha=16,
-        target_modules=['c_attn'],  # LoRA on the attention weights
+        target_modules=["c_attn"],  # LoRA on the attention weights
         lora_dropout=0.1,
-        bias='none',
+        bias="none",
     )
     lora_model = get_peft_model(model, config)
     return lora_model
@@ -167,9 +171,9 @@ lora_model.eval()
 
 lora_la = Laplace(
     lora_model,
-    likelihood='classification',
-    subset_of_weights='all',
-    hessian_structure='diag',
+    likelihood="classification",
+    subset_of_weights="all",
+    hessian_structure="diag",
     backend=AsdlGGN,
 )
 
@@ -199,7 +203,7 @@ subnetwork_indices = subnetwork_mask.select()
 
 # Example 2: specify the layers that define the subnetwork
 from laplace.utils import ModuleNameSubnetMask
-subnetwork_mask = ModuleNameSubnetMask(model, module_names=['layer.1', 'layer.3'])
+subnetwork_mask = ModuleNameSubnetMask(model, module_names=["layer.1", "layer.3"])
 subnetwork_mask.select()
 subnetwork_indices = subnetwork_mask.indices
 
@@ -208,9 +212,9 @@ import torch
 subnetwork_indices = torch.tensor([0, 4, 11, 42, 123, 2021])
 
 # Define and fit subnetwork LA using the specified subnetwork indices
-la = Laplace(model, 'classification',
-             subset_of_weights='subnetwork',
-             hessian_structure='full',
+la = Laplace(model, "classification",
+             subset_of_weights="subnetwork",
+             hessian_structure="full",
              subnetwork_indices=subnetwork_indices)
 la.fit(train_loader)
 ```
@@ -220,21 +224,21 @@ la.fit(train_loader)
 As with plain `torch`, we support to ways to serialize data.
 
 One is the familiar `state_dict` approach. Here you need to save and re-create
-both `model` and  `Laplace`. Use this for long-term storage of models and
+both `model` and `Laplace`. Use this for long-term storage of models and
 sharing of a fitted `Laplace` instance.
 
 ```py
 # Save model and Laplace instance
-torch.save(model.state_dict(), 'model_state_dict.bin')
-torch.save(la.state_dict(), 'la_state_dict.bin')
+torch.save(model.state_dict(), "model_state_dict.bin")
+torch.save(la.state_dict(), "la_state_dict.bin")
 
 # Load serialized data
 model2 = MyModel(...)
-model2.load_state_dict(torch.load('model_state_dict.bin'))
-la2 = Laplace(model2, 'classification',
-              subset_of_weights='all',
-              hessian_structure='diag')
-la2.load_state_dict(torch.load('la_state_dict.bin'))
+model2.load_state_dict(torch.load("model_state_dict.bin"))
+la2 = Laplace(model2, "classification",
+              subset_of_weights="all",
+              hessian_structure="diag")
+la2.load_state_dict(torch.load("la_state_dict.bin"))
 ```
 
 The second approach is to save the whole `Laplace` object, including
@@ -246,10 +250,10 @@ Use this for quick save-load cycles during experiments, say.
 
 ```py
 # Save Laplace, including la.model
-torch.save(la, 'la.pt')
+torch.save(la, "la.pt")
 
 # Load both
-torch.load('la.pt')
+torch.load("la.pt")
 ```
 
 Some Laplace variants such as `LLLaplace` might have trouble being serialized
@@ -260,22 +264,23 @@ using the default `pickle` module, which `torch.save()` and `torch.load()` use
 ```py
 import dill
 
-torch.save(la, 'la.pt', pickle_module=dill)
+torch.save(la, "la.pt", pickle_module=dill)
 ```
 
 With both methods, you are free to switch devices, for instance when you
 trained on a GPU but want to run predictions on CPU. In this case, use
 
 ```py
-torch.load(..., map_location='cpu')
+torch.load(..., map_location="cpu")
 ```
 
 ## Structure
+
 The laplace package consists of two main components:
 
-1. The subclasses of [`laplace.BaseLaplace`](https://github.com/AlexImmer/Laplace/blob/main/laplace/baselaplace.py) that implement different sparsity structures: different subsets of weights (`'all'`, `'subnetwork'` and `'last_layer'`) and different structures of the Hessian approximation (`'full'`, `'kron'`, `'lowrank'` and `'diag'`). This results in _nine_ currently available options: `laplace.FullLaplace`, `laplace.KronLaplace`, `laplace.DiagLaplace`, the corresponding last-layer variations `laplace.FullLLLaplace`, `laplace.KronLLLaplace`,  and `laplace.DiagLLLaplace` (which are all subclasses of [`laplace.LLLaplace`](https://github.com/AlexImmer/Laplace/blob/main/laplace/lllaplace.py)), [`laplace.SubnetLaplace`](https://github.com/AlexImmer/Laplace/blob/main/laplace/subnetlaplace.py) (which only supports `'full'` and `'diag'` Hessian approximations) and `laplace.LowRankLaplace` (which only supports inference over `'all'` weights). All of these can be conveniently accessed via the [`laplace.Laplace`](https://github.com/AlexImmer/Laplace/blob/main/laplace/laplace.py) function.
+1. The subclasses of [`laplace.BaseLaplace`](https://github.com/AlexImmer/Laplace/blob/main/laplace/baselaplace.py) that implement different sparsity structures: different subsets of weights (`"all"`, `"subnetwork"` and `"last_layer"`) and different structures of the Hessian approximation (`"full"`, `"kron"`, `"lowrank"` and `"diag"`). This results in _nine_ currently available options: `laplace.FullLaplace`, `laplace.KronLaplace`, `laplace.DiagLaplace`, the corresponding last-layer variations `laplace.FullLLLaplace`, `laplace.KronLLLaplace`, and `laplace.DiagLLLaplace` (which are all subclasses of [`laplace.LLLaplace`](https://github.com/AlexImmer/Laplace/blob/main/laplace/lllaplace.py)), [`laplace.SubnetLaplace`](https://github.com/AlexImmer/Laplace/blob/main/laplace/subnetlaplace.py) (which only supports `"full"` and `"diag"` Hessian approximations) and `laplace.LowRankLaplace` (which only supports inference over `"all"` weights). All of these can be conveniently accessed via the [`laplace.Laplace`](https://github.com/AlexImmer/Laplace/blob/main/laplace/laplace.py) function.
 2. The backends in [`laplace.curvature`](https://github.com/AlexImmer/Laplace/blob/main/laplace/curvature/) which provide access to Hessian approximations of
-the corresponding sparsity structures, for example, the diagonal GGN.
+   the corresponding sparsity structures, for example, the diagonal GGN.
 
 Additionally, the package provides utilities for
 decomposing a neural network into feature extractor and last layer for `LLLaplace` subclasses ([`laplace.utils.feature_extractor`](https://github.com/AlexImmer/Laplace/blob/main/laplace/utils/feature_extractor.py))
@@ -287,6 +292,7 @@ Automatic subnetwork selection strategies include: uniformly at random (`laplace
 In addition to that, subnetworks can also be specified manually, by listing the names of either the model parameters (`ParamNameSubnetMask`) or modules (`ModuleNameSubnetMask`) to perform Laplace inference over.
 
 ## Extendability
+
 To extend the laplace package, new `BaseLaplace` subclasses can be designed, for example,
 Laplace with a block-diagonal Hessian structure.
 One can also implement custom subnetwork selection strategies as new subclasses of `SubnetMask`.
@@ -317,14 +323,14 @@ pdoc --http 0.0.0.0:8080 laplace --template-dir template
 
 This package relies on various improvements to the Laplace approximation for neural networks, which was originally due to MacKay [1]. Please consider citing the respective papers if you use any of their proposed methods via our laplace library.
 
-- [1] MacKay, DJC. [*A Practical Bayesian Framework for Backpropagation Networks*](https://authors.library.caltech.edu/13793/). Neural Computation 1992.
-- [2] Gibbs, M. N. [*Bayesian Gaussian Processes for Regression and Classification*](https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.147.1130&rep=rep1&type=pdf). PhD Thesis 1997.
-- [3] Snoek, J., Rippel, O., Swersky, K., Kiros, R., Satish, N., Sundaram, N., Patwary, M., Prabhat, M., Adams, R. [*Scalable Bayesian Optimization Using Deep Neural Networks*](https://arxiv.org/abs/1502.05700). ICML 2015.
-- [4] Ritter, H., Botev, A., Barber, D. [*A Scalable Laplace Approximation for Neural Networks*](https://openreview.net/forum?id=Skdvd2xAZ). ICLR 2018.
-- [5] Foong, A. Y., Li, Y., Hernández-Lobato, J. M., Turner, R. E. [*'In-Between' Uncertainty in Bayesian Neural Networks*](https://arxiv.org/abs/1906.11537). ICML UDL Workshop 2019.
-- [6] Khan, M. E., Immer, A., Abedi, E., Korzepa, M. [*Approximate Inference Turns Deep Networks into Gaussian Processes*](https://arxiv.org/abs/1906.01930). NeurIPS 2019.
-- [7] Kristiadi, A., Hein, M., Hennig, P. [*Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks*](https://arxiv.org/abs/2002.10118). ICML 2020.
-- [8] Immer, A., Korzepa, M., Bauer, M. [*Improving predictions of Bayesian neural nets via local linearization*](https://arxiv.org/abs/2008.08400). AISTATS 2021.
-- [9] Sharma, A., Azizan, N., Pavone, M. [*Sketching Curvature for Efficient Out-of-Distribution Detection for Deep Neural Networks*](https://arxiv.org/abs/2102.12567). UAI 2021.
-- [10] Immer, A., Bauer, M., Fortuin, V., Rätsch, G., Khan, EM. [*Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning*](https://arxiv.org/abs/2104.04975). ICML 2021.
-- [11] Daxberger, E., Nalisnick, E., Allingham, JU., Antorán, J., Hernández-Lobato, JM. [*Bayesian Deep Learning via Subnetwork Inference*](https://arxiv.org/abs/2010.14689). ICML 2021.
+- [1] MacKay, DJC. [_A Practical Bayesian Framework for Backpropagation Networks_](https://authors.library.caltech.edu/13793/). Neural Computation 1992.
+- [2] Gibbs, M. N. [_Bayesian Gaussian Processes for Regression and Classification_](https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.147.1130&rep=rep1&type=pdf). PhD Thesis 1997.
+- [3] Snoek, J., Rippel, O., Swersky, K., Kiros, R., Satish, N., Sundaram, N., Patwary, M., Prabhat, M., Adams, R. [_Scalable Bayesian Optimization Using Deep Neural Networks_](https://arxiv.org/abs/1502.05700). ICML 2015.
+- [4] Ritter, H., Botev, A., Barber, D. [_A Scalable Laplace Approximation for Neural Networks_](https://openreview.net/forum?id=Skdvd2xAZ). ICLR 2018.
+- [5] Foong, A. Y., Li, Y., Hernández-Lobato, J. M., Turner, R. E. [_'In-Between' Uncertainty in Bayesian Neural Networks_](https://arxiv.org/abs/1906.11537). ICML UDL Workshop 2019.
+- [6] Khan, M. E., Immer, A., Abedi, E., Korzepa, M. [_Approximate Inference Turns Deep Networks into Gaussian Processes_](https://arxiv.org/abs/1906.01930). NeurIPS 2019.
+- [7] Kristiadi, A., Hein, M., Hennig, P. [_Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks_](https://arxiv.org/abs/2002.10118). ICML 2020.
+- [8] Immer, A., Korzepa, M., Bauer, M. [_Improving predictions of Bayesian neural nets via local linearization_](https://arxiv.org/abs/2008.08400). AISTATS 2021.
+- [9] Sharma, A., Azizan, N., Pavone, M. [_Sketching Curvature for Efficient Out-of-Distribution Detection for Deep Neural Networks_](https://arxiv.org/abs/2102.12567). UAI 2021.
+- [10] Immer, A., Bauer, M., Fortuin, V., Rätsch, G., Khan, EM. [_Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning_](https://arxiv.org/abs/2104.04975). ICML 2021.
+- [11] Daxberger, E., Nalisnick, E., Allingham, JU., Antorán, J., Hernández-Lobato, JM. [_Bayesian Deep Learning via Subnetwork Inference_](https://arxiv.org/abs/2010.14689). ICML 2021.
diff --git a/examples/bayesopt_example.py b/examples/bayesopt_example.py
index efe5eec..ff3fd3f 100644
--- a/examples/bayesopt_example.py
+++ b/examples/bayesopt_example.py
@@ -1,41 +1,28 @@
 import warnings
-warnings.filterwarnings('ignore')
+
+warnings.filterwarnings("ignore")
+
+import argparse
+import sys
+from typing import Any, List, Optional
 
 import numpy as np
-import functorch as ft
 import torch
-from torch import nn, optim
-from torch import distributions as dists
-from torch.nn import functional as F
-from gpytorch import distributions as gdists
 import torch.utils.data as data_utils
-from botorch.models.model import Model
-from botorch.posteriors.gpytorch import GPyTorchPosterior 
 import tqdm
-
-from laplace import Laplace
-from laplace.curvature import AsdlGGN, BackPackGGN, AsdlHessian
-
-import argparse, sys
-
-from botorch.models.gp_regression import SingleTaskGP
-from botorch.test_functions import Ackley, Branin
 from botorch.acquisition.analytic import ExpectedImprovement, UpperConfidenceBound
 from botorch.acquisition.monte_carlo import qExpectedImprovement
+from botorch.models.gp_regression import SingleTaskGP
+from botorch.models.model import Model
 from botorch.optim.optimize import optimize_acqf
+from botorch.posteriors.gpytorch import GPyTorchPosterior
+from botorch.test_functions import Ackley, Branin
+from gpytorch import distributions as gdists
+from torch import distributions as dists
+from torch import nn, optim
+from torch.nn import functional as F
 
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Hashable,
-    List,
-    Mapping,
-    Optional,
-    TYPE_CHECKING,
-    TypeVar,
-    Union,
-)
+from laplace import Laplace
 
 
 class LaplaceBNN(Model):
@@ -45,13 +32,13 @@ class LaplaceBNN(Model):
     """
 
     def __init__(
-            self,
-            train_X: torch.Tensor,
-            train_Y: torch.Tensor,
-            bnn: Laplace = None,
-            likelihood: str = 'regression',
-            batch_size: int = 1024):
-
+        self,
+        train_X: torch.Tensor,
+        train_Y: torch.Tensor,
+        bnn: Laplace = None,
+        likelihood: str = "regression",
+        batch_size: int = 1024,
+    ):
         super().__init__()
 
         self.train_X = train_X
@@ -63,14 +50,13 @@ def __init__(
             nn.ReLU(),
             nn.Linear(50, 50),
             nn.ReLU(),
-            nn.Linear(50, train_Y.shape[-1])
+            nn.Linear(50, train_Y.shape[-1]),
         )
         self.bnn = bnn
 
         if self.bnn is None:
             self._train_model(self._get_train_loader())
 
-
     def posterior(
         self,
         X: torch.Tensor,
@@ -84,7 +70,7 @@ def posterior(
         """
         # Transform to `(batch_shape*q, d)`
         B, Q, D = X.shape
-        X = X.reshape(B*Q, D)
+        X = X.reshape(B * Q, D)
 
         # Posterior predictive distribution
         # mean_y is (batch_shape*q, k); cov_y is (batch_shape*q*k, batch_shape*q*k)
@@ -92,13 +78,13 @@ def posterior(
 
         # Mean in `(batch_shape, q*k)`
         K = self.num_outputs
-        mean_y = mean_y.reshape(B, Q*K)
+        mean_y = mean_y.reshape(B, Q * K)
 
         # Cov is `(batch_shape, q*k, q*k)`
-        cov_y += 1e-4*torch.eye(B*Q*K)
+        cov_y += 1e-4 * torch.eye(B * Q * K)
         cov_y = cov_y.reshape(B, Q, K, B, Q, K)
-        cov_y = torch.einsum('bqkbrl->bqkrl', cov_y)  # (B, Q, K, Q, K)
-        cov_y = cov_y.reshape(B, Q*K, Q*K)
+        cov_y = torch.einsum("bqkbrl->bqkrl", cov_y)  # (B, Q, K, Q, K)
+        cov_y = cov_y.reshape(B, Q * K, Q * K)
 
         dist = gdists.MultivariateNormal(mean_y, covariance_matrix=cov_y)
         post_pred = GPyTorchPosterior(dist)
@@ -108,8 +94,9 @@ def posterior(
 
         return post_pred
 
-
-    def condition_on_observations(self, X: torch.Tensor, Y: torch.Tensor, **kwargs: Any) -> Model:
+    def condition_on_observations(
+        self, X: torch.Tensor, Y: torch.Tensor, **kwargs: Any
+    ) -> Model:
         self.train_X = torch.cat([self.train_X, X], dim=0)
         self.train_Y = torch.cat([self.train_Y, Y], dim=0)
 
@@ -118,8 +105,11 @@ def condition_on_observations(self, X: torch.Tensor, Y: torch.Tensor, **kwargs:
 
         return LaplaceBNN(
             # Added dataset & retrained BNN
-            self.train_X, self.train_Y, self.bnn,
-            self.likelihood, self.batch_size,
+            self.train_X,
+            self.train_Y,
+            self.bnn,
+            self.likelihood,
+            self.batch_size,
         )
 
     @property
@@ -133,7 +123,9 @@ def _train_model(self, train_loader):
         """
         n_epochs = 1000
         optimizer = optim.Adam(self.nn.parameters(), lr=1e-1, weight_decay=1e-3)
-        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, n_epochs*len(train_loader))
+        scheduler = optim.lr_scheduler.CosineAnnealingLR(
+            optimizer, n_epochs * len(train_loader)
+        )
         loss_func = nn.MSELoss()
 
         for i in range(n_epochs):
@@ -148,14 +140,15 @@ def _train_model(self, train_loader):
         self.nn.eval()
 
         self.bnn = Laplace(
-            self.nn, self.likelihood,
-            subset_of_weights='all', hessian_structure='kron',
-            enable_backprop=True
+            self.nn,
+            self.likelihood,
+            subset_of_weights="all",
+            hessian_structure="kron",
+            enable_backprop=True,
         )
         self.bnn.fit(train_loader)
         self.bnn.optimize_prior_precision(n_steps=50)
 
-
     def _get_prediction(self, test_x: torch.Tensor, joint=True, use_test_loader=False):
         """
         Batched Laplace prediction.
@@ -167,7 +160,7 @@ def _get_prediction(self, test_x: torch.Tensor, joint=True, use_test_loader=Fals
             Tensor of size `(batch_shape, k)`
         """
         if self.bnn is None:
-            print('Train your model first before making prediction!')
+            print("Train your model first before making prediction!")
             sys.exit(1)
 
         if not use_test_loader:
@@ -175,7 +168,7 @@ def _get_prediction(self, test_x: torch.Tensor, joint=True, use_test_loader=Fals
         else:
             test_loader = data_utils.DataLoader(
                 data_utils.TensorDataset(test_x, torch.zeros_like(test_x)),
-                batch_size=256
+                batch_size=256,
             )
 
             mean_y, cov_y = [], []
@@ -190,73 +183,78 @@ def _get_prediction(self, test_x: torch.Tensor, joint=True, use_test_loader=Fals
 
         return mean_y, cov_y
 
-
     def _get_train_loader(self):
         return data_utils.DataLoader(
             data_utils.TensorDataset(self.train_X, self.train_Y),
-            batch_size=self.batch_size, shuffle=True
+            batch_size=self.batch_size,
+            shuffle=True,
         )
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--test_func', choices=['ackley', 'branin'], default='branin')
-    parser.add_argument('--acqf', choices=['EI', 'UCB', 'qEI'], default='EI')
-    parser.add_argument('--init_data', type=int, default=20)
-    parser.add_argument('--exp_len', type=int, default=500)
-    parser.add_argument('--randseed', type=int, default=1)
+    parser.add_argument("--test_func", choices=["ackley", "branin"], default="branin")
+    parser.add_argument("--acqf", choices=["EI", "UCB", "qEI"], default="EI")
+    parser.add_argument("--init_data", type=int, default=20)
+    parser.add_argument("--exp_len", type=int, default=500)
+    parser.add_argument("--randseed", type=int, default=1)
     args = parser.parse_args()
 
     np.random.seed(args.randseed)
     torch.set_default_dtype(torch.float64)
     torch.manual_seed(args.randseed)
 
-    if args.test_func == 'ackley':
+    if args.test_func == "ackley":
         true_f = Ackley(dim=2)
         bounds = torch.tensor([[-32.768, 32.768], [-32.768, 32.768]]).T.double()
-    elif args.test_func == 'branin':
+    elif args.test_func == "branin":
         true_f = Branin()
         bounds = torch.tensor([[-5, 10], [0, 15]]).T.double()
     else:
-        print('Invalid test function!')
+        print("Invalid test function!")
         sys.exit(1)
 
     print()
-    print(f'Test Function: {args.test_func}')
-    print('-------------------------------------')
+    print(f"Test Function: {args.test_func}")
+    print("-------------------------------------")
     print()
 
     train_data_points = 20
 
-    train_x = torch.cat([
-        dists.Uniform(*bounds.T[i]).sample((train_data_points, 1))
-        for i in range(2)  # for each dimension
-    ], dim=1)
+    train_x = torch.cat(
+        [
+            dists.Uniform(*bounds.T[i]).sample((train_data_points, 1))
+            for i in range(2)  # for each dimension
+        ],
+        dim=1,
+    )
     train_y = true_f(train_x).reshape(-1, 1)
 
-    test_x = torch.cat([
-        dists.Uniform(*bounds.T[i]).sample((10000, 1))
-        for i in range(2)  # for each dimension
-    ], dim=1) 
+    test_x = torch.cat(
+        [
+            dists.Uniform(*bounds.T[i]).sample((10000, 1))
+            for i in range(2)  # for each dimension
+        ],
+        dim=1,
+    )
     test_y = true_f(test_x)
 
     models = {
-        'RandomSearch': None, 
-        'BNN-LA': LaplaceBNN(train_x, train_y),
-        'GP': SingleTaskGP(train_x, train_y),
+        "RandomSearch": None,
+        "BNN-LA": LaplaceBNN(train_x, train_y),
+        "GP": SingleTaskGP(train_x, train_y),
     }
 
-
     def evaluate_model(model_name, model):
-        if model_name == 'GP':
+        if model_name == "GP":
             pred = model.posterior(test_x).mean.squeeze()
-        elif model_name == 'BNN-LA':
+        elif model_name == "BNN-LA":
             pred, _ = model._get_prediction(test_x, use_test_loader=True, joint=False)
         else:
             return -1
 
         return F.mse_loss(pred, test_y).squeeze().item()
 
-
     for model_name, model in models.items():
         np.random.seed(args.randseed)
         torch.set_default_dtype(torch.float64)
@@ -266,43 +264,46 @@ def evaluate_model(model_name, model):
         trace_best_y = []
         pbar = tqdm.trange(args.exp_len)
         pbar.set_description(
-            f'[{model_name}, MSE = {evaluate_model(model_name, model):.3f}; Best f(x) = {best_y:.3f}]'
+            f"[{model_name}, MSE = {evaluate_model(model_name, model):.3f}; Best f(x) = {best_y:.3f}]"
         )
 
         for i in pbar:
-            if model_name == 'RandomSearch':
-                new_x = torch.cat([
-                    dists.Uniform(*bounds.T[i]).sample((1, 1))
-                    for i in range(len(bounds))  # for each dimension
-                ], dim=1).squeeze()
+            if model_name == "RandomSearch":
+                new_x = torch.cat(
+                    [
+                        dists.Uniform(*bounds.T[i]).sample((1, 1))
+                        for i in range(len(bounds))  # for each dimension
+                    ],
+                    dim=1,
+                ).squeeze()
             else:
-                if args.acqf == 'EI':
+                if args.acqf == "EI":
                     acq_f = ExpectedImprovement(model, best_f=best_y, maximize=False)
-                elif args.acqf == 'UCB':
+                elif args.acqf == "UCB":
                     acq_f = UpperConfidenceBound(model, beta=0.2, maximize=False)
-                elif args.acqf == 'qEI':
+                elif args.acqf == "qEI":
                     acq_f = qExpectedImprovement(model, best_f=best_y, maximize=False)
                 else:
-                    print('Invalid acquisition function!')
+                    print("Invalid acquisition function!")
                     sys.exit(1)
 
                 # Get a proposal for new x
                 new_x, val = optimize_acqf(
                     acq_f,
                     bounds=bounds,
-                    q=1 if args.acqf not in ['qEI'] else 5,
+                    q=1 if args.acqf not in ["qEI"] else 5,
                     num_restarts=10,
-                    raw_samples=20
+                    raw_samples=20,
                 )
 
             if len(new_x.shape) == 1:
                 new_x = new_x.unsqueeze(0)
-            
+
             # Evaluate the objective on the proposed x
             new_y = true_f(new_x).unsqueeze(-1)  # (q, 1)
 
             # Update posterior
-            if model_name != 'RandomSearch':
+            if model_name != "RandomSearch":
                 model = model.condition_on_observations(new_x, new_y)
 
             # Update the current best y
@@ -311,5 +312,5 @@ def evaluate_model(model_name, model):
 
             trace_best_y.append(best_y)
             pbar.set_description(
-                f'[{model_name}, MSE = {evaluate_model(model_name, model):.3f}; Best f(x) = {best_y:.3f}, curr f(x) = {new_y.min().item():.3f}]'
+                f"[{model_name}, MSE = {evaluate_model(model_name, model):.3f}; Best f(x) = {best_y:.3f}, curr f(x) = {new_y.min().item():.3f}]"
             )
diff --git a/examples/calibration_example.md b/examples/calibration_example.md
index d0a47d4..3d56c60 100644
--- a/examples/calibration_example.md
+++ b/examples/calibration_example.md
@@ -1,6 +1,6 @@
-## Full example: *post-hoc* Laplace on a large image classifier
+## Full example: _post-hoc_ Laplace on a large image classifier
 
-An advantage of the Laplace approximation over variational Bayes and Markov Chain Monte Carlo methods is its *post-hoc* nature. That means we can apply LA on (almost) any *pre-trained* neural network. In this example, we will see how we can apply the last-layer LA on a deep WideResNet model, trained on CIFAR-10.
+An advantage of the Laplace approximation over variational Bayes and Markov Chain Monte Carlo methods is its _post-hoc_ nature. That means we can apply LA on (almost) any _pre-trained_ neural network. In this example, we will see how we can apply the last-layer LA on a deep WideResNet model, trained on CIFAR-10.
 
 #### Data loading
 
@@ -32,18 +32,18 @@ targets = torch.cat([y for x, y in test_loader], dim=0).numpy()
 
 Next, we will load a pre-trained WideResNet-16-4 model. Note that a GPU with CUDA support is needed for this example.
 
-``` python
+```python
 # The model is a standard WideResNet 16-4
 # Taken as is from https://github.com/hendrycks/outlier-exposure
 model = wrn.WideResNet(16, 4, num_classes=10).cuda().eval()
 
 util.download_pretrained_model()
-model.load_state_dict(torch.load('./temp/CIFAR10_plain.pt'))
+model.load_state_dict(torch.load("./temp/CIFAR10_plain.pt"))
 ```
 
 To simplify the downstream tasks, we will use the following helper function to make predictions. It simply iterates through all minibatches and obtains the predictive probabilities of the CIFAR-10 classes.
 
-``` python
+```python
 @torch.no_grad()
 def predict(dataloader, model, laplace=False):
     py = []
@@ -63,13 +63,13 @@ We are now ready to see how calibrated is the model. The metrics we use are the
 
 First, let us inspect the MAP model. We shall use the [`netcal`](https://github.com/fabiankueppers/calibration-framework) library to easily compute the ECE.
 
-``` python
+```python
 probs_map = predict(test_loader, model, laplace=False)
 acc_map = (probs_map.argmax(-1) == targets).float().mean()
 ece_map = ECE(bins=15).measure(probs_map.numpy(), targets.numpy())
 nll_map = -dists.Categorical(probs_map).log_prob(targets).mean()
 
-print(f'[MAP] Acc.: {acc_map:.1%}; ECE: {ece_map:.1%}; NLL: {nll_map:.3}')
+print(f"[MAP] Acc.: {acc_map:.1%}; ECE: {ece_map:.1%}; NLL: {nll_map:.3}")
 ```
 
 Running this snippet, we would get:
@@ -80,26 +80,26 @@ Running this snippet, we would get:
 
 ### The calibration of Laplace
 
-Now we inspect the benefit of the LA. Let us apply the simple last-layer LA model, and optimize the prior precision hyperparameter using a *post-hoc* marginal likelihood maximization.
+Now we inspect the benefit of the LA. Let us apply the simple last-layer LA model, and optimize the prior precision hyperparameter using a _post-hoc_ marginal likelihood maximization.
 
-``` python
+```python
 # Laplace
-la = Laplace(model, 'classification',
-             subset_of_weights='last_layer',
-             hessian_structure='kron')
+la = Laplace(model, "classification",
+             subset_of_weights="last_layer",
+             hessian_structure="kron")
 la.fit(train_loader)
-la.optimize_prior_precision(method='marglik')
+la.optimize_prior_precision(method="marglik")
 ```
 
 Then, we are ready to see how well does LA improves the calibration of the MAP model:
 
-``` python
+```python
 probs_laplace = predict(test_loader, la, laplace=True)
 acc_laplace = (probs_laplace.argmax(-1) == targets).float().mean()
 ece_laplace = ECE(bins=15).measure(probs_laplace.numpy(), targets.numpy())
 nll_laplace = -dists.Categorical(probs_laplace).log_prob(targets).mean()
 
-print(f'[Laplace] Acc.: {acc_laplace:.1%}; ECE: {ece_laplace:.1%}; NLL: {nll_laplace:.3}')
+print(f"[Laplace] Acc.: {acc_laplace:.1%}; ECE: {ece_laplace:.1%}; NLL: {nll_laplace:.3}")
 ```
 
 Running this snippet, we obtain:
diff --git a/examples/calibration_example.py b/examples/calibration_example.py
index 3e97ed9..1eea9c0 100644
--- a/examples/calibration_example.py
+++ b/examples/calibration_example.py
@@ -1,18 +1,17 @@
 import warnings
 
-warnings.simplefilter('ignore', UserWarning)
-
-import torch  # noqa: E402
-import torch.distributions as dists  # noqa: E402
-import numpy as np  # noqa: E402
-import helper.wideresnet as wrn  # noqa: E402
-import helper.dataloaders as dl  # noqa: E402
-from helper import util  # noqa: E402
+warnings.simplefilter("ignore", UserWarning)
+
+import helper.dataloaders as dl
+import helper.wideresnet as wrn
+import numpy as np
+import torch
+import torch.distributions as dists
+from helper import util
 from netcal.metrics import ECE  # noqa: E402
 
 from laplace import Laplace  # noqa: E402
 
-
 np.random.seed(7777)
 torch.manual_seed(7777)
 torch.backends.cudnn.deterministic = True
@@ -27,7 +26,7 @@
 model = wrn.WideResNet(16, 4, num_classes=10).cuda().eval()
 
 util.download_pretrained_model()
-model.load_state_dict(torch.load('./temp/CIFAR10_plain.pt'))
+model.load_state_dict(torch.load("./temp/CIFAR10_plain.pt"))
 
 
 @torch.no_grad()
@@ -48,14 +47,14 @@ def predict(dataloader, model, laplace=False):
 ece_map = ECE(bins=15).measure(probs_map.numpy(), targets.numpy())
 nll_map = -dists.Categorical(probs_map).log_prob(targets).mean()
 
-print(f'[MAP] Acc.: {acc_map:.1%}; ECE: {ece_map:.1%}; NLL: {nll_map:.3}')
+print(f"[MAP] Acc.: {acc_map:.1%}; ECE: {ece_map:.1%}; NLL: {nll_map:.3}")
 
 # Laplace
 la = Laplace(
-    model, 'classification', subset_of_weights='last_layer', hessian_structure='kron'
+    model, "classification", subset_of_weights="last_layer", hessian_structure="kron"
 )
 la.fit(train_loader)
-la.optimize_prior_precision(method='marglik')
+la.optimize_prior_precision(method="marglik")
 
 probs_laplace = predict(test_loader, la, laplace=True)
 acc_laplace = (probs_laplace.argmax(-1) == targets).float().mean()
@@ -63,5 +62,5 @@ def predict(dataloader, model, laplace=False):
 nll_laplace = -dists.Categorical(probs_laplace).log_prob(targets).mean()
 
 print(
-    f'[Laplace] Acc.: {acc_laplace:.1%}; ECE: {ece_laplace:.1%}; NLL: {nll_laplace:.3}'
+    f"[Laplace] Acc.: {acc_laplace:.1%}; ECE: {ece_laplace:.1%}; NLL: {nll_laplace:.3}"
 )
diff --git a/examples/helper/dataloaders.py b/examples/helper/dataloaders.py
index 4107d59..a4fa21d 100644
--- a/examples/helper/dataloaders.py
+++ b/examples/helper/dataloaders.py
@@ -1,13 +1,12 @@
 # Taken from https://github.com/AlexMeinke/certified-certain-uncertainty
 import torch
-from torchvision import datasets, transforms
 import torch.utils.data as data_utils
-
+from torchvision import datasets, transforms
 
 train_batch_size = 128
 test_batch_size = 100
 
-path = './temp/'
+path = "./temp/"
 
 
 def CIFAR10(train=True, batch_size=None, augm_flag=True):
@@ -21,7 +20,7 @@ def CIFAR10(train=True, batch_size=None, augm_flag=True):
     transform_train = transforms.Compose(
         [
             transforms.RandomHorizontalFlip(),
-            transforms.RandomCrop(32, padding=4, padding_mode='reflect'),
+            transforms.RandomCrop(32, padding=4, padding_mode="reflect"),
         ]
         + transform_base
     )
diff --git a/examples/helper/util.py b/examples/helper/util.py
index 2f714ed..6f1005d 100644
--- a/examples/helper/util.py
+++ b/examples/helper/util.py
@@ -6,45 +6,45 @@
 
 def download_pretrained_model():
     # Download pre-trained model if necessary
-    if not os.path.isfile('CIFAR10_plain.pt'):
-        if not os.path.exists('./temp'):
-            os.makedirs('./temp')
+    if not os.path.isfile("CIFAR10_plain.pt"):
+        if not os.path.exists("./temp"):
+            os.makedirs("./temp")
 
         urllib.request.urlretrieve(
-            'https://nc.mlcloud.uni-tuebingen.de/index.php/s/2PBDYDsiotN76mq/download',
-            './temp/CIFAR10_plain.pt',
+            "https://nc.mlcloud.uni-tuebingen.de/index.php/s/2PBDYDsiotN76mq/download",
+            "./temp/CIFAR10_plain.pt",
         )
 
 
 def plot_regression(
-    X_train, y_train, X_test, f_test, y_std, plot=True, file_name='regression_example'
+    X_train, y_train, X_test, f_test, y_std, plot=True, file_name="regression_example"
 ):
     fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, sharey=True, figsize=(4.5, 2.8))
-    ax1.set_title('MAP')
-    ax1.scatter(X_train.flatten(), y_train.flatten(), alpha=0.3, color='tab:orange')
-    ax1.plot(X_test, f_test, color='black', label='$f_{MAP}$')
+    ax1.set_title("MAP")
+    ax1.scatter(X_train.flatten(), y_train.flatten(), alpha=0.3, color="tab:orange")
+    ax1.plot(X_test, f_test, color="black", label=r"$f_{MAP}$")
     ax1.legend()
 
-    ax2.set_title('LA')
-    ax2.scatter(X_train.flatten(), y_train.flatten(), alpha=0.3, color='tab:orange')
-    ax2.plot(X_test, f_test, label=r'$\mathbb{E}[f]$')
+    ax2.set_title("LA")
+    ax2.scatter(X_train.flatten(), y_train.flatten(), alpha=0.3, color="tab:orange")
+    ax2.plot(X_test, f_test, label=r"$\mathbb{E}[f]$")
     ax2.fill_between(
         X_test,
         f_test - y_std * 2,
         f_test + y_std * 2,
         alpha=0.3,
-        color='tab:blue',
-        label=r'$2\sqrt{\mathbb{V}\,[y]}$',
+        color="tab:blue",
+        label=r"$2\sqrt{\mathbb{V}\,[y]}$",
     )
     ax2.legend()
     ax1.set_ylim([-4, 6])
     ax1.set_xlim([X_test.min(), X_test.max()])
     ax2.set_xlim([X_test.min(), X_test.max()])
-    ax1.set_ylabel('$y$')
-    ax1.set_xlabel('$x$')
-    ax2.set_xlabel('$x$')
+    ax1.set_ylabel("$y$")
+    ax1.set_xlabel("$x$")
+    ax2.set_xlabel("$x$")
     plt.tight_layout()
     if plot:
         plt.show()
     else:
-        plt.savefig(f'docs/{file_name}.png')
+        plt.savefig(f"docs/{file_name}.png")
diff --git a/examples/helper/wideresnet.py b/examples/helper/wideresnet.py
index 387d596..3022b3a 100644
--- a/examples/helper/wideresnet.py
+++ b/examples/helper/wideresnet.py
@@ -1,4 +1,5 @@
 import math
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -10,14 +11,29 @@ def __init__(self, in_planes, out_planes, stride, dropRate=0.0):
 
         self.bn1 = nn.BatchNorm2d(in_planes)
         self.relu1 = nn.ReLU(inplace=True)
-        self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.conv1 = nn.Conv2d(
+            in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
+        )
 
         self.bn2 = nn.BatchNorm2d(out_planes)
         self.relu2 = nn.ReLU(inplace=True)
-        self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv2 = nn.Conv2d(
+            out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False
+        )
         self.droprate = dropRate
-        self.equalInOut = (in_planes == out_planes)
-        self.convShortcut = (not self.equalInOut) and nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False) or None
+        self.equalInOut = in_planes == out_planes
+        self.convShortcut = (
+            (not self.equalInOut)
+            and nn.Conv2d(
+                in_planes,
+                out_planes,
+                kernel_size=1,
+                stride=stride,
+                padding=0,
+                bias=False,
+            )
+            or None
+        )
 
     def forward(self, x):
         if not self.equalInOut:
@@ -44,13 +60,22 @@ def forward(self, x):
 class NetworkBlock(nn.Module):
     def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0):
         super(NetworkBlock, self).__init__()
-        self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate)
+        self.layer = self._make_layer(
+            block, in_planes, out_planes, nb_layers, stride, dropRate
+        )
 
     def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate):
         layers = []
 
         for i in range(nb_layers):
-            layers.append(block(i == 0 and in_planes or out_planes, out_planes, i == 0 and stride or 1, dropRate))
+            layers.append(
+                block(
+                    i == 0 and in_planes or out_planes,
+                    out_planes,
+                    i == 0 and stride or 1,
+                    dropRate,
+                )
+            )
 
         return nn.Sequential(*layers)
 
@@ -59,17 +84,26 @@ def forward(self, x):
 
 
 class WideResNet(nn.Module):
-
-    def __init__(self, depth, widen_factor, num_classes, num_channel=3, dropRate=0.3, feature_extractor=False):
+    def __init__(
+        self,
+        depth,
+        widen_factor,
+        num_classes,
+        num_channel=3,
+        dropRate=0.3,
+        feature_extractor=False,
+    ):
         super(WideResNet, self).__init__()
 
         nChannels = [16, 16 * widen_factor, 32 * widen_factor, 64 * widen_factor]
-        assert ((depth - 4) % 6 == 0)
+        assert (depth - 4) % 6 == 0
         n = (depth - 4) // 6
         block = BasicBlock
 
         # 1st conv before any network block
-        self.conv1 = nn.Conv2d(num_channel, nChannels[0], kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv1 = nn.Conv2d(
+            num_channel, nChannels[0], kernel_size=3, stride=1, padding=1, bias=False
+        )
         self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate)
         self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, dropRate)
         self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, dropRate)
@@ -84,7 +118,7 @@ def __init__(self, depth, widen_factor, num_classes, num_channel=3, dropRate=0.3
         for m in self.modules():
             if isinstance(m, nn.Conv2d):
                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
-                m.weight.data.normal_(0, math.sqrt(2. / n))
+                m.weight.data.normal_(0, math.sqrt(2.0 / n))
             elif isinstance(m, nn.BatchNorm2d):
                 m.weight.data.fill_(1)
                 m.bias.data.zero_()
@@ -99,7 +133,6 @@ def forward(self, x):
 
         return self.fc(out)
 
-
     def features(self, x):
         out = self.conv1(x)
         out = self.block1(out)
diff --git a/examples/huggingface_example.md b/examples/huggingface_example.md
index 6f54a07..1257c52 100644
--- a/examples/huggingface_example.md
+++ b/examples/huggingface_example.md
@@ -24,8 +24,8 @@ from laplace import Laplace
 import logging
 import warnings
 
-logging.basicConfig(level='ERROR')
-warnings.filterwarnings('ignore')
+logging.basicConfig(level="ERROR")
+warnings.filterwarnings("ignore")
 
 from transformers import ( # noqa: E402
     GPT2Config,
@@ -46,29 +46,29 @@ numpy.random.seed(0)
 Next, we create a toy dataset. You can use any HF datasets or your own, of course.
 
 ```python
-tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 tokenizer.pad_token_id = tokenizer.eos_token_id
 
 data = [
-    {'text': 'Today is hot, but I will manage!!!!', 'label': 1},
-    {'text': 'Tomorrow is cold', 'label': 0},
-    {'text': 'Carpe diem', 'label': 1},
-    {'text': 'Tempus fugit', 'label': 1},
+    {"text": "Today is hot, but I will manage!!!!", "label": 1},
+    {"text": "Tomorrow is cold", "label": 0},
+    {"text": "Carpe diem", "label": 1},
+    {"text": "Tempus fugit", "label": 1},
 ]
 dataset = Dataset.from_list(data)
 
 def tokenize(row):
-    return tokenizer(row['text'])
+    return tokenizer(row["text"])
 
-dataset = dataset.map(tokenize, remove_columns=['text'])
-dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
+dataset = dataset.map(tokenize, remove_columns=["text"])
+dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
 dataloader = data_utils.DataLoader(
 dataset, batch_size=100, collate_fn=DataCollatorWithPadding(tokenizer)
 )
 
 data = next(iter(dataloader))
 print(
-    f'Huggingface data defaults to UserDict, which is a MutableMapping? {isinstance(data, UserDict)}'
+    f"Huggingface data defaults to UserDict, which is a MutableMapping? {isinstance(data, UserDict)}"
 )
 for k, v in data.items():
     print(k, v.shape)
@@ -101,15 +101,15 @@ class MyGPT2(nn.Module):
 
     def __init__(self, tokenizer: PreTrainedTokenizer) -> None:
         super().__init__()
-        config = GPT2Config.from_pretrained('gpt2')
+        config = GPT2Config.from_pretrained("gpt2")
         config.pad_token_id = tokenizer.pad_token_id
         config.num_labels = 2
         self.hf_model = GPT2ForSequenceClassification.from_pretrained(
-            'gpt2', config=config
+            "gpt2", config=config
         )
 
     def forward(self, data: MutableMapping) -> torch.Tensor:
-        '''
+        """
         Custom forward function. Handles things like moving the
         input tensor to the correct device inside.
 
@@ -120,19 +120,19 @@ class MyGPT2(nn.Module):
 
         Returns:
             logits: An `(batch_size, n_classes)`-sized tensor of logits.
-        '''
+        """
         device = next(self.parameters()).device
-        input_ids = data['input_ids'].to(device)
-        attn_mask = data['attention_mask'].to(device)
+        input_ids = data["input_ids"].to(device)
+        attn_mask = data["attention_mask"].to(device)
         output_dict = self.hf_model(input_ids=input_ids, attention_mask=attn_mask)
         return output_dict.logits
 
 model = MyGPT2(tokenizer)
 ```
 
-Now, let's apply Laplace. Let's do a last-layer Laplace first. Notice that we add 
-an argument `feature_reduction` there. This is because Huggingface models reduce the 
-logits and [not the features](https://github.com/huggingface/transformers/blob/a98c41798cf6ed99e1ff17e3792d6e06a2ff2ff3/src/transformers/models/gpt2/modeling_gpt2.py#L1678-L1704). 
+Now, let's apply Laplace. Let's do a last-layer Laplace first. Notice that we add
+an argument `feature_reduction` there. This is because Huggingface models reduce the
+logits and [not the features](https://github.com/huggingface/transformers/blob/a98c41798cf6ed99e1ff17e3792d6e06a2ff2ff3/src/transformers/models/gpt2/modeling_gpt2.py#L1678-L1704).
 
 ```python
 model = MyGPT2(tokenizer)
@@ -140,29 +140,30 @@ model.eval()
 
 la = Laplace(
     model,
-    likelihood='classification',
-    subset_of_weights='last_layer',
-    hessian_structure='full',
+    likelihood="classification",
+    subset_of_weights="last_layer",
+    hessian_structure="full",
     # This must reflect faithfully the reduction technique used in the model
     # Otherwise, correctness is not guaranteed
-    feature_reduction=FeatureReduction.PICK_LAST,
+    feature_reduction="pick_last",
 )
 la.fit(dataloader)
 la.optimize_prior_precision()
 
 X_test = next(iter(dataloader))
-print(f'[Last-layer Laplace] The predictive tensor is of shape: {la(X_test).shape}.')
+print(f"[Last-layer Laplace] The predictive tensor is of shape: {la(X_test).shape}.")
 ```
 
 Here's the output:
+
 ```
 [Last-layer Laplace] The predictive tensor is of shape: torch.Size([4, 2]).
 ```
 
 ## Subnetwork Laplace
 
-Also, we can do the same thing by switching off the gradients of all layers except the 
-top layer. Laplace will automatically only compute the Hessian (and Jacobians) of the 
+Also, we can do the same thing by switching off the gradients of all layers except the
+top layer. Laplace will automatically only compute the Hessian (and Jacobians) of the
 parameters in which `requires_grad` is `True`.
 
 Notice that you can "mix-and-match" this gradient switching. You can do a subnetwork Laplace
@@ -182,15 +183,15 @@ for p in model.hf_model.score.parameters():
 la = Laplace(
     model,
     # Will only hit the last-layer since it's the only one that is grad-enabled
-    likelihood='classification',
-    subset_of_weights='all',
-    hessian_structure='diag',
+    likelihood="classification",
+    subset_of_weights="all",
+    hessian_structure="diag",
 )
 la.fit(dataloader)
 la.optimize_prior_precision()
 
 X_test = next(iter(dataloader))
-print(f'[Subnetwork Laplace] The predictive tensor is of shape: {la(X_test).shape}.')
+print(f"[Subnetwork Laplace] The predictive tensor is of shape: {la(X_test).shape}.")
 ```
 
 Here are the outputs to validate that Laplace works:
@@ -211,9 +212,9 @@ def get_lora_model():
         config = LoraConfig(
         r=4,
         lora_alpha=16,
-        target_modules=['c_attn'], # LoRA on the attention weights
+        target_modules=["c_attn"], # LoRA on the attention weights
         lora_dropout=0.1,
-        bias='none',
+        bias="none",
     )
     lora_model = get_peft_model(model, config)
     return lora_model
@@ -226,14 +227,14 @@ lora_model.eval()
 
 lora_la = Laplace(
     lora_model,
-    likelihood='classification',
-    subset_of_weights='all',
-    hessian_structure='kron',
+    likelihood="classification",
+    subset_of_weights="all",
+    hessian_structure="kron",
 )
 lora_la.fit(dataloader)
 
 X_test = next(iter(dataloader))
-print(f'[LoRA-LLM] The predictive tensor is of shape: {lora_la(X_test).shape}.')
+print(f"[LoRA-LLM] The predictive tensor is of shape: {lora_la(X_test).shape}.")
 ```
 
 Here is the output, as expected:
@@ -261,7 +262,7 @@ Then, any backend and any hessian factorization can be used in this case.
 
 Otherwise, if you must use dict-based inputs, choose the following backends:
 
-- `CurvlinopsGGN` for `hessian_factorization = {'kron', 'diag'}`
-- `CurvlinopsEF` for `hessian_factorization = {'kron'}`
-- `AsdlGGN` for `hessian_factorization = {'kron', 'diag'}`
-- `AsdlEF` for `hessian_factorization = {'kron', 'diag'}`
+- `CurvlinopsGGN` for `hessian_factorization = {"kron", "diag"}`
+- `CurvlinopsEF` for `hessian_factorization = {"kron"}`
+- `AsdlGGN` for `hessian_factorization = {"kron", "diag"}`
+- `AsdlEF` for `hessian_factorization = {"kron", "diag"}`
diff --git a/examples/huggingface_example.py b/examples/huggingface_example.py
index ec9cf9a..6c20947 100644
--- a/examples/huggingface_example.py
+++ b/examples/huggingface_example.py
@@ -6,50 +6,51 @@
 import numpy
 import torch
 import torch.utils.data as data_utils
+from datasets import Dataset
+from peft import LoraConfig, get_peft_model
 from torch import nn
+from transformers import (
+    DataCollatorWithPadding,
+    GPT2Config,
+    GPT2ForSequenceClassification,
+    GPT2Tokenizer,
+    PreTrainedTokenizer,
+)
 
 from laplace import Laplace
 
-from laplace.utils.feature_extractor import FeatureReduction
-
-logging.basicConfig(level='ERROR')
-warnings.filterwarnings('ignore')
-
-from datasets import Dataset  # noqa: E402
-from peft import LoraConfig, get_peft_model  # noqa: E402
-from transformers import (DataCollatorWithPadding, GPT2Config,  # noqa: E402
-                          GPT2ForSequenceClassification, GPT2Tokenizer,
-                          PreTrainedTokenizer)
+logging.basicConfig(level="ERROR")
+warnings.filterwarnings("ignore")
 
 # make deterministic
 torch.manual_seed(0)
 numpy.random.seed(0)
 
-tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 tokenizer.pad_token_id = tokenizer.eos_token_id
 
 data = [
-    {'text': 'Today is hot, but I will manage!!!!', 'label': 1},
-    {'text': 'Tomorrow is cold', 'label': 0},
-    {'text': 'Carpe diem', 'label': 1},
-    {'text': 'Tempus fugit', 'label': 1},
+    {"text": "Today is hot, but I will manage!!!!", "label": 1},
+    {"text": "Tomorrow is cold", "label": 0},
+    {"text": "Carpe diem", "label": 1},
+    {"text": "Tempus fugit", "label": 1},
 ]
 dataset = Dataset.from_list(data)
 
 
 def tokenize(row):
-    return tokenizer(row['text'])
+    return tokenizer(row["text"])
 
 
-dataset = dataset.map(tokenize, remove_columns=['text'])
-dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
+dataset = dataset.map(tokenize, remove_columns=["text"])
+dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
 dataloader = data_utils.DataLoader(
     dataset, batch_size=100, collate_fn=DataCollatorWithPadding(tokenizer)
 )
 
 data = next(iter(dataloader))
 print(
-    f'Huggingface data defaults to UserDict, which is a MutableMapping? {isinstance(data, UserDict)}'
+    f"Huggingface data defaults to UserDict, which is a MutableMapping? {isinstance(data, UserDict)}"
 )
 for k, v in data.items():
     print(k, v.shape)
@@ -66,11 +67,11 @@ class MyGPT2(nn.Module):
 
     def __init__(self, tokenizer: PreTrainedTokenizer) -> None:
         super().__init__()
-        config = GPT2Config.from_pretrained('gpt2')
+        config = GPT2Config.from_pretrained("gpt2")
         config.pad_token_id = tokenizer.pad_token_id
         config.num_labels = 2
         self.hf_model = GPT2ForSequenceClassification.from_pretrained(
-            'gpt2', config=config
+            "gpt2", config=config
         )
 
     def forward(self, data: MutableMapping) -> torch.Tensor:
@@ -87,8 +88,8 @@ def forward(self, data: MutableMapping) -> torch.Tensor:
             logits: An `(batch_size, n_classes)`-sized tensor of logits.
         """
         device = next(self.parameters()).device
-        input_ids = data['input_ids'].to(device)
-        attn_mask = data['attention_mask'].to(device)
+        input_ids = data["input_ids"].to(device)
+        attn_mask = data["attention_mask"].to(device)
         output_dict = self.hf_model(input_ids=input_ids, attention_mask=attn_mask)
         return output_dict.logits
 
@@ -100,18 +101,18 @@ def forward(self, data: MutableMapping) -> torch.Tensor:
 
 la = Laplace(
     model,
-    likelihood='classification',
-    subset_of_weights='last_layer',
-    hessian_structure='full',
+    likelihood="classification",
+    subset_of_weights="last_layer",
+    hessian_structure="full",
     # This must reflect faithfully the reduction technique used in the model
     # Otherwise, correctness is not guaranteed
-    feature_reduction=FeatureReduction.PICK_LAST,
+    feature_reduction="pick_last",
 )
 la.fit(dataloader)
 la.optimize_prior_precision()
 
 X_test = next(iter(dataloader))
-print(f'[Last-layer Laplace] The predictive tensor is of shape: {la(X_test).shape}.')
+print(f"[Last-layer Laplace] The predictive tensor is of shape: {la(X_test).shape}.")
 
 del model
 del la
@@ -131,16 +132,16 @@ def forward(self, data: MutableMapping) -> torch.Tensor:
 
 la = Laplace(
     model,
-    likelihood='classification',
+    likelihood="classification",
     # Will only hit the last-layer since it's the only one that is grad-enabled
-    subset_of_weights='all',
-    hessian_structure='diag',
+    subset_of_weights="all",
+    hessian_structure="diag",
 )
 la.fit(dataloader)
 la.optimize_prior_precision()
 
 X_test = next(iter(dataloader))
-print(f'[Subnetwork Laplace] The predictive tensor is of shape: {la(X_test).shape}.')
+print(f"[Subnetwork Laplace] The predictive tensor is of shape: {la(X_test).shape}.")
 
 del model
 del la
@@ -155,9 +156,9 @@ def get_lora_model():
     config = LoraConfig(
         r=4,
         lora_alpha=16,
-        target_modules=['c_attn'],  # LoRA on the attention weights
+        target_modules=["c_attn"],  # LoRA on the attention weights
         lora_dropout=0.1,
-        bias='none',
+        bias="none",
     )
     lora_model = get_peft_model(model, config)
     return lora_model
@@ -170,11 +171,11 @@ def get_lora_model():
 
 lora_la = Laplace(
     lora_model,
-    likelihood='classification',
-    subset_of_weights='all',
-    hessian_structure='kron',
+    likelihood="classification",
+    subset_of_weights="all",
+    hessian_structure="kron",
 )
 # lora_la.fit(dataloader)
 
 X_test = next(iter(dataloader))
-print(f'[LoRA-LLM] The predictive tensor is of shape: {lora_la(X_test).shape}.')
+print(f"[LoRA-LLM] The predictive tensor is of shape: {lora_la(X_test).shape}.")
diff --git a/examples/regression_example.md b/examples/regression_example.md
index 6ce4049..77dcc2b 100644
--- a/examples/regression_example.md
+++ b/examples/regression_example.md
@@ -1,10 +1,12 @@
 ## Full example: Optimization of the marginal likelihood and prediction
 
 ### Sinusoidal toy data
-We show how the marginal likelihood can be used after training a MAP network on a simple sinusoidal regression task. 
+
+We show how the marginal likelihood can be used after training a MAP network on a simple sinusoidal regression task.
 Subsequently, we use the optimized LA to predict which provides uncertainty on top of the MAP prediction.
 We also show how the `marglik_training` utility method can be used to jointly train the MAP and hyperparameters.
 First, we set up the training data for the problem with observation noise \\(\\sigma=0.3\\):
+
 ```python
 from laplace.baselaplace import FullLaplace
 from laplace.curvature.backpack import BackPackGGN
@@ -23,9 +25,11 @@ X_train, y_train, train_loader, X_test = get_sinusoid_example(sigma_noise=0.3)
 ```
 
 ### Training a MAP
+
 We now use `pytorch` to train a neural network with single hidden layer and Tanh activation.
 The trained neural network will be our MAP estimate.
 This is standard so nothing new here, yet:
+
 ```python
 # create and train MAP model
 def get_model():
@@ -46,15 +50,16 @@ for i in range(n_epochs):
 ```
 
 ### Fitting and optimizing the Laplace approximation using empirical Bayes
+
 With the MAP-trained model at hand, we can estimate the prior precision and observation noise
 using empirical Bayes after training.
-The `Laplace` method is called to construct a LA for `'regression'` with `'all'` weights.
-As default `Laplace` returns a Kronecker factored LA, we use `'full'` instead on this small example.
+The `Laplace` method is called to construct a LA for `"regression"` with `"all"` weights.
+As default `Laplace` returns a Kronecker factored LA, we use `"full"` instead on this small example.
 We fit the LA to the training data and initialize `log_prior` and `log_sigma`.
 Using Adam, we minimize the negative log marginal likelihood for `n_epochs`.
 
 ```python
-la = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='full')
+la = Laplace(model, "regression", subset_of_weights="all", hessian_structure="full")
 la.fit(train_loader)
 log_prior, log_sigma = torch.ones(1, requires_grad=True), torch.ones(1, requires_grad=True)
 hyper_optimizer = torch.optim.Adam([log_prior, log_sigma], lr=1e-1)
@@ -70,9 +75,10 @@ without the need for any validation data.
 The resulting prior precision is \\(\\delta \\approx 0.10\\).
 
 ### Bayesian predictive
+
 Here, we compare the MAP prediction to the obtained LA prediction.
 For LA, we have a closed-form predictive distribution on the output \\(f\\) which is a Gaussian
-\\(\\mathcal{N}(f(x;\\theta_{MAP}), \\mathbb{V}[f] + \\sigma^2)\\):
+\\(\\mathcal{N}(f(x;\\theta\_{MAP}), \\mathbb{V}[f] + \\sigma^2)\\):
 
 ```python
 x = X_test.flatten().cpu().numpy()
@@ -85,13 +91,14 @@ plot_regression(X_train, y_train, x, f_mu, pred_std)
 ```
 
 .. image:: regression_example.png
-   :align: center
+:align: center
 
 In comparison to the MAP, the predictive shows useful uncertainties.
 When our MAP is over or underfit, the Laplace approximation cannot fix this anymore.
 In this case, joint optimization of MAP and marginal likelihood can be useful.
 
 ### Jointly optimize MAP and hyperparameters using online empirical Bayes
+
 We provide a utility method `marglik_training` that implements the algorithm proposed in [1].
 The method optimizes the neural network and the hyperparameters in an interleaved way
 and returns an optimally regularized LA.
@@ -100,9 +107,9 @@ Below, we use this method and plot the corresponding predictive uncertainties ag
 ```python
 model = get_model()
 la, model, margliks, losses = marglik_training(
-    model=model, train_loader=train_loader, likelihood='regression',
-    hessian_structure='full', backend=BackPackGGN, n_epochs=n_epochs, 
-    optimizer_kwargs={'lr': 1e-2}, prior_structure='scalar'
+    model=model, train_loader=train_loader, likelihood="regression",
+    hessian_structure="full", backend=BackPackGGN, n_epochs=n_epochs,
+    optimizer_kwargs={"lr": 1e-2}, prior_structure="scalar"
 )
 
 f_mu, f_var = la(X_test)
@@ -114,4 +121,4 @@ plot_regression(X_train, y_train, x, f_mu, pred_std)
 ```
 
 .. image:: regression_example_online.png
-   :align: center
+:align: center
diff --git a/examples/regression_example.py b/examples/regression_example.py
index 0158df2..d7ddf44 100644
--- a/examples/regression_example.py
+++ b/examples/regression_example.py
@@ -1,13 +1,10 @@
-from laplace.baselaplace import FullLaplace
-from laplace.curvature.backpack import BackPackGGN
 import numpy as np
 import torch
-
-from laplace import Laplace, marglik_training
-
 from helper.dataloaders import get_sinusoid_example
 from helper.util import plot_regression
 
+from laplace import Laplace, marglik_training
+from laplace.curvature.backpack import BackPackGGN
 
 n_epochs = 1000
 torch.manual_seed(711)
@@ -36,7 +33,7 @@ def get_model():
         loss.backward()
         optimizer.step()
 
-la = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='full')
+la = Laplace(model, "regression", subset_of_weights="all", hessian_structure="full")
 la.fit(train_loader)
 log_prior, log_sigma = (
     torch.ones(1, requires_grad=True),
@@ -51,15 +48,15 @@ def get_model():
 
 # Serialization for fitted quantities
 state_dict = la.state_dict()
-torch.save(state_dict, 'state_dict.bin')
+torch.save(state_dict, "state_dict.bin")
 
-la = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='full')
+la = Laplace(model, "regression", subset_of_weights="all", hessian_structure="full")
 # Load serialized, fitted quantities
-la.load_state_dict(torch.load('state_dict.bin'))
+la.load_state_dict(torch.load("state_dict.bin"))
 
 print(
-    f'sigma={la.sigma_noise.item():.2f}',
-    f'prior precision={la.prior_precision.item():.2f}',
+    f"sigma={la.sigma_noise.item():.2f}",
+    f"prior precision={la.prior_precision.item():.2f}",
 )
 
 x = X_test.flatten().cpu().numpy()
@@ -82,7 +79,7 @@ def get_model():
 pred_std = np.sqrt(f_sigma**2 + la.sigma_noise.item() ** 2)
 
 plot_regression(
-    X_train, y_train, x, f_mu, pred_std, file_name='regression_example', plot=False
+    X_train, y_train, x, f_mu, pred_std, file_name="regression_example", plot=False
 )
 
 # alternatively, optimize parameters and hyperparameters of the prior jointly
@@ -90,17 +87,17 @@ def get_model():
 la, model, margliks, losses = marglik_training(
     model=model,
     train_loader=train_loader,
-    likelihood='regression',
-    hessian_structure='full',
+    likelihood="regression",
+    hessian_structure="full",
     backend=BackPackGGN,
     n_epochs=n_epochs,
-    optimizer_kwargs={'lr': 1e-2},
-    prior_structure='scalar',
+    optimizer_kwargs={"lr": 1e-2},
+    prior_structure="scalar",
 )
 
 print(
-    f'sigma={la.sigma_noise.item():.2f}',
-    f'prior precision={la.prior_precision.numpy()}',
+    f"sigma={la.sigma_noise.item():.2f}",
+    f"prior precision={la.prior_precision.numpy()}",
 )
 
 f_mu, f_var = la(X_test)
@@ -113,6 +110,6 @@ def get_model():
     x,
     f_mu,
     pred_std,
-    file_name='regression_example_online',
+    file_name="regression_example_online",
     plot=False,
 )
diff --git a/examples/reward_modeling_example.md b/examples/reward_modeling_example.md
index 5f61914..a338572 100644
--- a/examples/reward_modeling_example.md
+++ b/examples/reward_modeling_example.md
@@ -1,14 +1,13 @@
 ## Full Example: Bayesian Bradley-Terry Reward Modeling
 
-The `laplace-torch` library can also be used to "Bayesianize" a pretrained Bradley-Terry 
-reward model, popular in large language models. See <http://arxiv.org/abs/2009.01325> 
+The `laplace-torch` library can also be used to "Bayesianize" a pretrained Bradley-Terry
+reward model, popular in large language models. See <http://arxiv.org/abs/2009.01325>
 for a primer in reward modeling.
 
 First order of business, let's define our comparison dataset. We will use the `datasets`
 library from Huggingface to handle the data.
 
-
-``` python
+```python
 import numpy as np
 import torch
 from torch import nn, optim
@@ -22,8 +21,8 @@ from laplace import Laplace
 import logging
 import warnings
 
-logging.basicConfig(level='ERROR')
-warnings.filterwarnings('ignore')
+logging.basicConfig(level="ERROR")
+warnings.filterwarnings("ignore")
 
 # make deterministic
 torch.manual_seed(0)
@@ -33,16 +32,16 @@ np.random.seed(0)
 # Pairwise comparison dataset. The label indicates which `x0` or `x1` is preferred.
 data_dict = [
     {
-        'x0': torch.randn(3),
-        'x1': torch.randn(3),
-        'label': torch.randint(2, size=(1,)).item(),
+        "x0": torch.randn(3),
+        "x1": torch.randn(3),
+        "label": torch.randint(2, size=(1,)).item(),
     }
     for _ in range(10)
 ]
 dataset = Dataset.from_list(data_dict)
 ```
 
-Now, let's define the reward model.  During training, it assumes that `x` is a tensor 
+Now, let's define the reward model. During training, it assumes that `x` is a tensor
 of shape `(batch_size, 2, dim)`, which is a concatenation of `x0` and `x1` above.
 The second dimension of size 2 is preserved through the forward pass, resulting in
 a logit tensor of shape `(batch_size, 2)` (the network itself is single-output).
@@ -53,10 +52,10 @@ prefer to use the dict-like inputs as in Huggingface LLM models, this can also b
 Simply combine what you have learned from this example with the Huggingface LLM example
 provided in this library.
 
-During testing, this model behaves like a standard single-output regression 
+During testing, this model behaves like a standard single-output regression
 model.
 
-``` python
+```python
 class SimpleRewardModel(nn.Module):
     """A simple reward model, compatible with the Bradley-Terry likelihood.
     """
@@ -92,26 +91,26 @@ class SimpleRewardModel(nn.Module):
             return logits
 ```
 
-To fulfill the 3D tensor requirement, we need to preprocess the dict-based dataset. 
+To fulfill the 3D tensor requirement, we need to preprocess the dict-based dataset.
 
-``` python
+```python
 # Preprocess to coalesce x0 and x1 into a single array/tensor
 def append_x0_x1(row):
     # The tensor values above are automatically casted as lists by `Dataset`
-    row['x'] = np.stack([row['x0'], row['x1']])  # (2, dim)
+    row["x"] = np.stack([row["x0"], row["x1"]])  # (2, dim)
     return row
 
 
-tensor_dataset = dataset.map(append_x0_x1, remove_columns=['x0', 'x1'])
-tensor_dataset.set_format(type='torch', columns=['x', 'label'])
+tensor_dataset = dataset.map(append_x0_x1, remove_columns=["x0", "x1"])
+tensor_dataset.set_format(type="torch", columns=["x", "label"])
 tensor_dataloader = data_utils.DataLoader(
-    data_utils.TensorDataset(tensor_dataset['x'], tensor_dataset['label']), batch_size=3
+    data_utils.TensorDataset(tensor_dataset["x"], tensor_dataset["label"]), batch_size=3
 )
 ```
 
 Then, we can train as usual using the cross entropy loss.
 
-``` python
+```python
 reward_model = SimpleRewardModel()
 opt = optim.AdamW(reward_model.parameters(), weight_decay=1e-3)
 
@@ -127,26 +126,26 @@ for epoch in range(10):
 
 Applying Laplace to this model is a breeze. Simply state that the likelihood is `reward_modeling`.
 
-``` python
+```python
 # Laplace !!! Notice the likelihood !!!
 reward_model.eval()
-la = Laplace(reward_model, likelihood='reward_modeling', subset_of_weights='all')
+la = Laplace(reward_model, likelihood="reward_modeling", subset_of_weights="all")
 la.fit(tensor_dataloader)
 la.optimize_prior_precision()
 ```
 
 As we can see, during prediction, even though we train & fit Laplace using the cross entropy
 loss (i.e. classification), in test time, the model behaves like a regression model.
-So, you don't get probability vectors as outputs. Instead, you get two tensors 
+So, you don't get probability vectors as outputs. Instead, you get two tensors
 containing the predictive means and predictive variance.
 
-``` python
+```python
 x_test = torch.randn(5, 3)
 pred_mean, pred_var = la(x_test)
 print(
-    f'Input shape {tuple(x_test.shape)}, predictive mean of shape '
-    + f'{tuple(pred_mean.shape)}, predictive covariance of shape '
-    + f'{tuple(pred_var.shape)}'
+    f"Input shape {tuple(x_test.shape)}, predictive mean of shape "
+    + f"{tuple(pred_mean.shape)}, predictive covariance of shape "
+    + f"{tuple(pred_var.shape)}"
 )
 ```
 
diff --git a/examples/reward_modeling_example.py b/examples/reward_modeling_example.py
index 7e8b632..6c65054 100644
--- a/examples/reward_modeling_example.py
+++ b/examples/reward_modeling_example.py
@@ -1,18 +1,17 @@
+import logging
+import warnings
+
 import numpy as np
 import torch
-from torch import nn, optim
-from torch.nn import functional as F
 import torch.utils.data as data_utils
-
 from datasets import Dataset
+from torch import nn, optim
+from torch.nn import functional as F
 
 from laplace import Laplace
 
-import logging
-import warnings
-
-logging.basicConfig(level='ERROR')
-warnings.filterwarnings('ignore')
+logging.basicConfig(level="ERROR")
+warnings.filterwarnings("ignore")
 
 # make deterministic
 torch.manual_seed(0)
@@ -22,9 +21,9 @@
 # Pairwise comparison dataset. The label indicates which `x0` or `x1` is preferred.
 data_dict = [
     {
-        'x0': torch.randn(3),
-        'x1': torch.randn(3),
-        'label': torch.randint(2, size=(1,)).item(),
+        "x0": torch.randn(3),
+        "x1": torch.randn(3),
+        "label": torch.randint(2, size=(1,)).item(),
     }
     for _ in range(10)
 ]
@@ -77,14 +76,14 @@ def forward(self, x):
 # Preprocess to coalesce x0 and x1 into a single array/tensor
 def append_x0_x1(row):
     # The tensor values above are automatically casted as lists by `Dataset`
-    row['x'] = np.stack([row['x0'], row['x1']])  # (2, dim)
+    row["x"] = np.stack([row["x0"], row["x1"]])  # (2, dim)
     return row
 
 
-tensor_dataset = dataset.map(append_x0_x1, remove_columns=['x0', 'x1'])
-tensor_dataset.set_format(type='torch', columns=['x', 'label'])
+tensor_dataset = dataset.map(append_x0_x1, remove_columns=["x0", "x1"])
+tensor_dataset.set_format(type="torch", columns=["x", "label"])
 tensor_dataloader = data_utils.DataLoader(
-    data_utils.TensorDataset(tensor_dataset['x'], tensor_dataset['label']), batch_size=3
+    data_utils.TensorDataset(tensor_dataset["x"], tensor_dataset["label"]), batch_size=3
 )
 
 reward_model = SimpleRewardModel()
@@ -101,14 +100,14 @@ def append_x0_x1(row):
 
 # Laplace !!! Notice the likelihood !!!
 reward_model.eval()
-la = Laplace(reward_model, likelihood='reward_modeling', subset_of_weights='all')
+la = Laplace(reward_model, likelihood="reward_modeling", subset_of_weights="all")
 la.fit(tensor_dataloader)
 la.optimize_prior_precision()
 
 x_test = torch.randn(5, 3)
 pred_mean, pred_var = la(x_test)
 print(
-    f'Input shape {tuple(x_test.shape)}, predictive mean of shape '
-    + f'{tuple(pred_mean.shape)}, predictive covariance of shape '
-    + f'{tuple(pred_var.shape)}'
+    f"Input shape {tuple(x_test.shape)}, predictive mean of shape "
+    + f"{tuple(pred_mean.shape)}, predictive covariance of shape "
+    + f"{tuple(pred_var.shape)}"
 )
diff --git a/laplace/__init__.py b/laplace/__init__.py
index 303306f..5446db2 100644
--- a/laplace/__init__.py
+++ b/laplace/__init__.py
@@ -6,36 +6,36 @@
 .. include:: ../examples/reward_modeling_example.md
 """
 
-REGRESSION = 'regression'
-CLASSIFICATION = 'classification'
+REGRESSION = "regression"
+CLASSIFICATION = "classification"
 
 from laplace.baselaplace import (
     BaseLaplace,
-    ParametricLaplace,
+    DiagLaplace,
     FullLaplace,
     KronLaplace,
-    DiagLaplace,
     LowRankLaplace,
+    ParametricLaplace,
 )
-from laplace.lllaplace import LLLaplace, FullLLLaplace, KronLLLaplace, DiagLLLaplace
-from laplace.subnetlaplace import SubnetLaplace, FullSubnetLaplace, DiagSubnetLaplace
 from laplace.laplace import Laplace
+from laplace.lllaplace import DiagLLLaplace, FullLLLaplace, KronLLLaplace, LLLaplace
 from laplace.marglik_training import marglik_training
+from laplace.subnetlaplace import DiagSubnetLaplace, FullSubnetLaplace, SubnetLaplace
 
 __all__ = [
-    'Laplace',  # direct access to all Laplace classes via unified interface
-    'BaseLaplace',
-    'ParametricLaplace',  # base-class and its (first-level) subclasses
-    'FullLaplace',
-    'KronLaplace',
-    'DiagLaplace',
-    'LowRankLaplace',  # all-weights
-    'LLLaplace',  # base-class last-layer
-    'FullLLLaplace',
-    'KronLLLaplace',
-    'DiagLLLaplace',  # last-layer
-    'SubnetLaplace',  # base-class subnetwork
-    'FullSubnetLaplace',
-    'DiagSubnetLaplace',  # subnetwork
-    'marglik_training',
+    "Laplace",  # direct access to all Laplace classes via unified interface
+    "BaseLaplace",
+    "ParametricLaplace",  # base-class and its (first-level) subclasses
+    "FullLaplace",
+    "KronLaplace",
+    "DiagLaplace",
+    "LowRankLaplace",  # all-weights
+    "LLLaplace",  # base-class last-layer
+    "FullLLLaplace",
+    "KronLLLaplace",
+    "DiagLLLaplace",  # last-layer
+    "SubnetLaplace",  # base-class subnetwork
+    "FullSubnetLaplace",
+    "DiagSubnetLaplace",  # subnetwork
+    "marglik_training",
 ]  # methods
diff --git a/laplace/baselaplace.py b/laplace/baselaplace.py
index e00ba8c..e3233bd 100644
--- a/laplace/baselaplace.py
+++ b/laplace/baselaplace.py
@@ -11,16 +11,22 @@
 from laplace.curvature import CurvlinopsGGN
 from laplace.curvature.asdfghjkl import AsdfghjklHessian
 from laplace.curvature.curvlinops import CurvlinopsEF
-from laplace.utils import (Kron, RunningNLLMetric, fix_prior_prec_structure,
-                           invsqrt_precision, normal_samples, validate)
+from laplace.utils import (
+    Kron,
+    RunningNLLMetric,
+    fix_prior_prec_structure,
+    invsqrt_precision,
+    normal_samples,
+    validate,
+)
 
 __all__ = [
-    'BaseLaplace',
-    'ParametricLaplace',
-    'FullLaplace',
-    'KronLaplace',
-    'DiagLaplace',
-    'LowRankLaplace',
+    "BaseLaplace",
+    "ParametricLaplace",
+    "FullLaplace",
+    "KronLaplace",
+    "DiagLaplace",
+    "LowRankLaplace",
 ]
 
 
@@ -76,14 +82,14 @@ def __init__(
         prior_mean=0.0,
         temperature=1.0,
         enable_backprop=False,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
         backend=None,
         backend_kwargs=None,
         asdl_fisher_kwargs=None,
     ):
-        if likelihood not in ['classification', 'regression', 'reward_modeling']:
-            raise ValueError(f'Invalid likelihood type {likelihood}')
+        if likelihood not in ["classification", "regression", "reward_modeling"]:
+            raise ValueError(f"Invalid likelihood type {likelihood}")
 
         self.model = model
 
@@ -100,13 +106,13 @@ def __init__(
         self.n_layers = len(self.params)
         self.prior_precision = prior_precision
         self.prior_mean = prior_mean
-        if sigma_noise != 1 and likelihood != 'regression':
-            raise ValueError('Sigma noise != 1 only available for regression.')
+        if sigma_noise != 1 and likelihood != "regression":
+            raise ValueError("Sigma noise != 1 only available for regression.")
 
-        self.reward_modeling = likelihood == 'reward_modeling'
+        self.reward_modeling = likelihood == "reward_modeling"
         if self.reward_modeling:
             # For fitting only. After it's done, self.likelihood = 'regression', see self.fit()
-            self.likelihood = 'classification'
+            self.likelihood = "classification"
         else:
             self.likelihood = likelihood
 
@@ -121,9 +127,13 @@ def __init__(
         if backend is None:
             backend = CurvlinopsGGN
         else:
-            if self.is_subset_params and 'backpack' in backend.__name__.lower():
+            if self.is_subset_params and (
+                "backpack" in backend.__name__.lower()
+                or "asdfghjkl" in backend.__name__.lower()
+            ):
                 raise ValueError(
-                    'If some grad are switched off, the BackPACK backend is not supported.'
+                    "If some grad are switched off, the BackPACK and Asdfghjkl backends"
+                    " are not supported."
                 )
 
         self._backend = None
@@ -175,7 +185,7 @@ def log_likelihood(self):
         log_likelihood : torch.Tensor
         """
         factor = -self._H_factor
-        if self.likelihood == 'regression':
+        if self.likelihood == "regression":
             # loss used is just MSE, need to add normalizer for gaussian likelihood
             c = (
                 self.n_data
@@ -195,12 +205,12 @@ def predictive(self, x, pred_type, link_approx, n_samples):
 
     def _check_jacobians(self, Js):
         if not isinstance(Js, torch.Tensor):
-            raise ValueError('Jacobians have to be torch.Tensor.')
+            raise ValueError("Jacobians have to be torch.Tensor.")
         if not Js.device == self._device:
-            raise ValueError('Jacobians need to be on the same device as Laplace.')
+            raise ValueError("Jacobians need to be on the same device as Laplace.")
         m, k, p = Js.size()
         if p != self.n_params:
-            raise ValueError('Invalid Jacobians shape for Laplace posterior approx.')
+            raise ValueError("Invalid Jacobians shape for Laplace posterior approx.")
 
     @property
     def prior_precision_diag(self):
@@ -228,7 +238,7 @@ def prior_precision_diag(self):
 
         else:
             raise ValueError(
-                'Mismatch of prior and model. Diagonal, scalar, or per-layer prior.'
+                "Mismatch of prior and model. Diagonal, scalar, or per-layer prior."
             )
 
     @property
@@ -244,12 +254,12 @@ def prior_mean(self, prior_mean):
                 self._prior_mean = prior_mean.reshape(-1).to(self._device)
             elif prior_mean.ndim == 1:
                 if len(prior_mean) not in [1, self.n_params]:
-                    raise ValueError('Invalid length of prior mean.')
+                    raise ValueError("Invalid length of prior mean.")
                 self._prior_mean = prior_mean
             else:
-                raise ValueError('Prior mean has too many dimensions!')
+                raise ValueError("Prior mean has too many dimensions!")
         else:
-            raise ValueError('Invalid argument type of prior mean.')
+            raise ValueError("Invalid argument type of prior mean.")
 
     @property
     def prior_precision(self):
@@ -267,32 +277,32 @@ def prior_precision(self, prior_precision):
             elif prior_precision.ndim == 1:
                 if len(prior_precision) not in [1, self.n_layers, self.n_params]:
                     raise ValueError(
-                        'Length of prior precision does not align with architecture.'
+                        "Length of prior precision does not align with architecture."
                     )
                 self._prior_precision = prior_precision.to(self._device)
             else:
                 raise ValueError(
-                    'Prior precision needs to be at most one-dimensional tensor.'
+                    "Prior precision needs to be at most one-dimensional tensor."
                 )
         else:
             raise ValueError(
-                'Prior precision either scalar or torch.Tensor up to 1-dim.'
+                "Prior precision either scalar or torch.Tensor up to 1-dim."
             )
 
     def optimize_prior_precision_base(
         self,
         pred_type,
-        method='marglik',
+        method="marglik",
         n_steps=100,
         lr=1e-1,
         init_prior_prec=1.0,
-        prior_structure='scalar',
+        prior_structure="scalar",
         val_loader=None,
         loss=None,
         log_prior_prec_min=-4,
         log_prior_prec_max=4,
         grid_size=100,
-        link_approx='probit',
+        link_approx="probit",
         n_samples=100,
         verbose=False,
         cv_loss_with_var=False,
@@ -346,9 +356,9 @@ def optimize_prior_precision_base(
             whether to show a progress bar; updated at every batch-Hessian computation.
             Useful for very large model and large amount of data, esp. when `subset_of_weights='all'`.
         """
-        if method == 'marglik':
+        if method == "marglik":
             self.prior_precision = init_prior_prec
-            if len(self.prior_precision) == 1 and prior_structure != 'scalar':
+            if len(self.prior_precision) == 1 and prior_structure != "scalar":
                 self.prior_precision = fix_prior_prec_structure(
                     self.prior_precision.item(),
                     prior_structure,
@@ -362,7 +372,7 @@ def optimize_prior_precision_base(
 
             if progress_bar:
                 pbar = tqdm.trange(n_steps)
-                pbar.set_description('[Optimizing marginal likelihood]')
+                pbar.set_description("[Optimizing marginal likelihood]")
             else:
                 pbar = range(n_steps)
 
@@ -375,16 +385,16 @@ def optimize_prior_precision_base(
                 neg_log_marglik.backward()
                 optimizer.step()
             self.prior_precision = log_prior_prec.detach().exp()
-        elif method == 'gridsearch':
+        elif method == "gridsearch":
             if val_loader is None:
-                raise ValueError('gridsearch requires a validation set DataLoader')
+                raise ValueError("gridsearch requires a validation set DataLoader")
 
             interval = torch.logspace(log_prior_prec_min, log_prior_prec_max, grid_size)
 
             if loss is None:
                 loss = (
                     tm.MeanSquaredError(num_outputs=self.n_outputs)
-                    if self.likelihood == 'regression'
+                    if self.likelihood == "regression"
                     else RunningNLLMetric()
                 )
 
@@ -399,9 +409,9 @@ def optimize_prior_precision_base(
                 progress_bar=progress_bar,
             )
         else:
-            raise ValueError('For now only marglik and gridsearch is implemented.')
+            raise ValueError("For now only marglik and gridsearch is implemented.")
         if verbose:
-            print(f'Optimized prior precision is {self.prior_precision}.')
+            print(f"Optimized prior precision is {self.prior_precision}.")
 
     def _gridsearch(
         self,
@@ -409,7 +419,7 @@ def _gridsearch(
         interval,
         val_loader,
         pred_type,
-        link_approx='probit',
+        link_approx="probit",
         n_samples=100,
         loss_with_var=False,
         progress_bar=False,
@@ -437,7 +447,7 @@ def _gridsearch(
 
             if progress_bar:
                 pbar.set_description(
-                    f'[Grid search | prior_prec: {prior_prec:.3e}, loss: {result:.3f}]'
+                    f"[Grid search | prior_prec: {prior_prec:.3e}, loss: {result:.3f}]"
                 )
 
             results.append(result)
@@ -458,13 +468,13 @@ def sigma_noise(self, sigma_noise):
                 self._sigma_noise = sigma_noise.to(self._device)
             elif sigma_noise.ndim == 1:
                 if len(sigma_noise) > 1:
-                    raise ValueError('Only homoscedastic output noise supported.')
+                    raise ValueError("Only homoscedastic output noise supported.")
                 self._sigma_noise = sigma_noise[0].to(self._device)
             else:
-                raise ValueError('Sigma noise needs to be scalar or 1-dimensional.')
+                raise ValueError("Sigma noise needs to be scalar or 1-dimensional.")
         else:
             raise ValueError(
-                'Invalid type: sigma noise needs to be torch.Tensor or scalar.'
+                "Invalid type: sigma noise needs to be torch.Tensor or scalar."
             )
 
     @property
@@ -506,8 +516,8 @@ def __init__(
         prior_mean=0.0,
         temperature=1.0,
         enable_backprop=False,
-        dict_key_x='inputs_id',
-        dict_key_y='labels',
+        dict_key_x="inputs_id",
+        dict_key_y="labels",
         backend=None,
         backend_kwargs=None,
         asdl_fisher_kwargs=None,
@@ -526,7 +536,7 @@ def __init__(
             backend_kwargs,
             asdl_fisher_kwargs,
         )
-        if not hasattr(self, 'H'):
+        if not hasattr(self, "H"):
             self._init_H()
             # posterior mean/mode
             self.mean = self.prior_mean
@@ -535,8 +545,8 @@ def _init_H(self):
         raise NotImplementedError
 
     def _check_H_init(self):
-        if getattr(self, 'H', None) is None:
-            raise AttributeError('Laplace not fitted. Run fit() first.')
+        if getattr(self, "H", None) is None:
+            raise AttributeError("Laplace not fitted. Run fit() first.")
 
     def fit(self, train_loader, override=True, progress_bar=False):
         """Fit the local Laplace approximation at the parameters of the model.
@@ -567,15 +577,15 @@ def fit(self, train_loader, override=True, progress_bar=False):
         data = next(iter(train_loader))
         with torch.no_grad():
             if isinstance(data, MutableMapping):  # To support Huggingface dataset
-                if 'backpack' in self._backend_cls.__name__.lower() or (
+                if "backpack" in self._backend_cls.__name__.lower() or (
                     isinstance(self, DiagLaplace) and self._backend_cls == CurvlinopsEF
                 ):
                     raise ValueError(
-                        'Currently DiagEF is not supported under CurvlinopsEF backend '
-                        + 'for custom models with non-tensor inputs '
-                        + '(https://github.com/pytorch/functorch/issues/159). Consider '
-                        + 'using AsdlEF backend instead. The same limitation applies '
-                        + 'to all BackPACK backend'
+                        "Currently DiagEF is not supported under CurvlinopsEF backend "
+                        + "for custom models with non-tensor inputs "
+                        + "(https://github.com/pytorch/functorch/issues/159). Consider "
+                        + "using AsdlEF backend instead. The same limitation applies "
+                        + "to all BackPACK backend"
                     )
 
                 out = self.model(data)
@@ -586,12 +596,12 @@ def fit(self, train_loader, override=True, progress_bar=False):
                 except (TypeError, AttributeError):
                     out = self.model(X.to(self._device))
         self.n_outputs = out.shape[-1]
-        setattr(self.model, 'output_size', self.n_outputs)
+        setattr(self.model, "output_size", self.n_outputs)
 
         N = len(train_loader.dataset)
         if progress_bar:
             pbar = tqdm.tqdm(train_loader)
-            pbar.set_description('[Computing Hessian]')
+            pbar.set_description("[Computing Hessian]")
         else:
             pbar = train_loader
 
@@ -716,8 +726,8 @@ def log_marginal_likelihood(self, prior_precision=None, sigma_noise=None):
 
         # update sigma_noise (useful when iterating on marglik)
         if sigma_noise is not None:
-            if self.likelihood != 'regression':
-                raise ValueError('Can only change sigma_noise for regression.')
+            if self.likelihood != "regression":
+                raise ValueError("Can only change sigma_noise for regression.")
             self.sigma_noise = sigma_noise
 
         return self.log_likelihood - 0.5 * (self.log_det_ratio + self.scatter)
@@ -725,9 +735,9 @@ def log_marginal_likelihood(self, prior_precision=None, sigma_noise=None):
     def __call__(
         self,
         x,
-        pred_type='glm',
+        pred_type="glm",
         joint=False,
-        link_approx='probit',
+        link_approx="probit",
         n_samples=100,
         diagonal_output=False,
         generator=None,
@@ -779,15 +789,15 @@ def __call__(
             For `likelihood='regression'` and `joint=True`, a tuple of torch.Tensor
             is returned with the mean and the predictive covariance.
         """
-        if pred_type not in ['glm', 'nn']:
-            raise ValueError('Only glm and nn supported as prediction types.')
+        if pred_type not in ["glm", "nn"]:
+            raise ValueError("Only glm and nn supported as prediction types.")
 
-        if link_approx not in ['mc', 'probit', 'bridge', 'bridge_norm']:
-            raise ValueError(f'Unsupported link approximation {link_approx}.')
+        if link_approx not in ["mc", "probit", "bridge", "bridge_norm"]:
+            raise ValueError(f"Unsupported link approximation {link_approx}.")
 
-        if pred_type == 'nn' and link_approx != 'mc':
+        if pred_type == "nn" and link_approx != "mc":
             raise ValueError(
-                'Only mc link approximation is supported for nn prediction type.'
+                "Only mc link approximation is supported for nn prediction type."
             )
 
         if generator is not None:
@@ -795,32 +805,32 @@ def __call__(
                 not isinstance(generator, torch.Generator)
                 or generator.device != x.device
             ):
-                raise ValueError('Invalid random generator (check type and device).')
+                raise ValueError("Invalid random generator (check type and device).")
 
         # For reward modeling, replace the likelihood to regression and override model state
-        if self.reward_modeling and self.likelihood == 'classification':
-            self.likelihood = 'regression'
+        if self.reward_modeling and self.likelihood == "classification":
+            self.likelihood = "regression"
             self.model.output_size = 1
 
-        if pred_type == 'glm':
+        if pred_type == "glm":
             f_mu, f_var = self._glm_predictive_distribution(
-                x, joint=joint and self.likelihood == 'regression'
+                x, joint=joint and self.likelihood == "regression"
             )
             # regression
-            if self.likelihood == 'regression':
+            if self.likelihood == "regression":
                 return f_mu, f_var
             # classification
-            if link_approx == 'mc':
+            if link_approx == "mc":
                 return self.predictive_samples(
                     x,
-                    pred_type='glm',
+                    pred_type="glm",
                     n_samples=n_samples,
                     diagonal_output=diagonal_output,
                 ).mean(dim=0)
-            elif link_approx == 'probit':
+            elif link_approx == "probit":
                 kappa = 1 / torch.sqrt(1.0 + np.pi / 8 * f_var.diagonal(dim1=1, dim2=2))
                 return torch.softmax(kappa * f_mu, dim=-1)
-            elif 'bridge' in link_approx:
+            elif "bridge" in link_approx:
                 # zero mean correction
                 f_mu -= (
                     f_var.sum(-1)
@@ -828,13 +838,13 @@ def __call__(
                     / f_var.sum(dim=(1, 2)).reshape(-1, 1)
                 )
                 f_var -= torch.einsum(
-                    'bi,bj->bij', f_var.sum(-1), f_var.sum(-2)
+                    "bi,bj->bij", f_var.sum(-1), f_var.sum(-2)
                 ) / f_var.sum(dim=(1, 2)).reshape(-1, 1, 1)
                 # Laplace Bridge
                 _, K = f_mu.size(0), f_mu.size(-1)
                 f_var_diag = torch.diagonal(f_var, dim1=1, dim2=2)
                 # optional: variance correction
-                if link_approx == 'bridge_norm':
+                if link_approx == "bridge_norm":
                     f_var_diag_mean = f_var_diag.mean(dim=1)
                     f_var_diag_mean /= torch.as_tensor(
                         [K / 2], device=self._device
@@ -845,14 +855,14 @@ def __call__(
                 alpha = (1 - 2 / K + f_mu.exp() / K**2 * sum_exp) / f_var_diag
                 return torch.nan_to_num(alpha / alpha.sum(dim=1).unsqueeze(-1), nan=1.0)
         else:
-            if self.likelihood == 'regression':
+            if self.likelihood == "regression":
                 samples = self._nn_predictive_samples(x, n_samples, **model_kwargs)
                 return samples.mean(dim=0), samples.var(dim=0)
             else:  # classification; the average is computed online
                 return self._nn_predictive_classification(x, n_samples, **model_kwargs)
 
     def predictive_samples(
-        self, x, pred_type='glm', n_samples=100, diagonal_output=False, generator=None
+        self, x, pred_type="glm", n_samples=100, diagonal_output=False, generator=None
     ):
         """Sample from the posterior predictive on input data `x`.
         Can be used, for example, for Thompson sampling.
@@ -882,10 +892,10 @@ def predictive_samples(
         samples : torch.Tensor
             samples `(n_samples, batch_size, output_shape)`
         """
-        if pred_type not in ['glm', 'nn']:
-            raise ValueError('Only glm and nn supported as prediction types.')
+        if pred_type not in ["glm", "nn"]:
+            raise ValueError("Only glm and nn supported as prediction types.")
 
-        if pred_type == 'glm':
+        if pred_type == "glm":
             f_mu, f_var = self._glm_predictive_distribution(x)
             assert f_var.shape == torch.Size(
                 [f_mu.shape[0], f_mu.shape[1], f_mu.shape[1]]
@@ -893,7 +903,7 @@ def predictive_samples(
             if diagonal_output:
                 f_var = torch.diagonal(f_var, dim1=1, dim2=2)
             f_samples = normal_samples(f_mu, f_var, n_samples, generator)
-            if self.likelihood == 'regression':
+            if self.likelihood == "regression":
                 return f_samples
             else:
                 return torch.softmax(f_samples, dim=-1)
@@ -903,7 +913,7 @@ def predictive_samples(
 
     @torch.enable_grad()
     def _glm_predictive_distribution(self, X, joint=False):
-        if 'backpack' in self._backend_cls.__name__.lower():
+        if "backpack" in self._backend_cls.__name__.lower():
             # BackPACK supports backprop through Jacobians, but it interferes with functorch
             Js, f_mu = self.backend.jacobians(X, enable_backprop=self.enable_backprop)
         else:
@@ -934,7 +944,7 @@ def _nn_predictive_samples(self, X, n_samples=100, generator=None, **model_kwarg
             fs.append(logits.detach() if not self.enable_backprop else logits)
         vector_to_parameters(self.mean, self.params)
         fs = torch.stack(fs)
-        if self.likelihood == 'classification':
+        if self.likelihood == "classification":
             fs = torch.softmax(fs, dim=-1)
         return fs
 
@@ -1008,24 +1018,24 @@ def sample(self, n_samples=100, generator=None):
 
     def optimize_prior_precision(
         self,
-        method='marglik',
-        pred_type='glm',
+        method="marglik",
+        pred_type="glm",
         n_steps=100,
         lr=1e-1,
         init_prior_prec=1.0,
-        prior_structure='scalar',
+        prior_structure="scalar",
         val_loader=None,
         loss=None,
         log_prior_prec_min=-4,
         log_prior_prec_max=4,
         grid_size=100,
-        link_approx='probit',
+        link_approx="probit",
         n_samples=100,
         verbose=False,
         cv_loss_with_var=False,
         progress_bar=False,
     ):
-        assert pred_type in ['glm', 'nn']
+        assert pred_type in ["glm", "nn"]
         self.optimize_prior_precision_base(
             pred_type,
             method,
@@ -1058,64 +1068,64 @@ def posterior_precision(self):
     def state_dict(self) -> dict:
         self._check_H_init()
         state_dict = {
-            'mean': self.mean,
-            'H': self.H,
-            'loss': self.loss,
-            'prior_mean': self.prior_mean,
-            'prior_precision': self.prior_precision,
-            'sigma_noise': self.sigma_noise,
-            'n_data': self.n_data,
-            'n_outputs': self.n_outputs,
-            'likelihood': self.likelihood,
-            'temperature': self.temperature,
-            'enable_backprop': self.enable_backprop,
-            'cls_name': self.__class__.__name__,
+            "mean": self.mean,
+            "H": self.H,
+            "loss": self.loss,
+            "prior_mean": self.prior_mean,
+            "prior_precision": self.prior_precision,
+            "sigma_noise": self.sigma_noise,
+            "n_data": self.n_data,
+            "n_outputs": self.n_outputs,
+            "likelihood": self.likelihood,
+            "temperature": self.temperature,
+            "enable_backprop": self.enable_backprop,
+            "cls_name": self.__class__.__name__,
         }
         return state_dict
 
     def load_state_dict(self, state_dict: dict):
         # Dealbreaker errors
-        if self.__class__.__name__ != state_dict['cls_name']:
+        if self.__class__.__name__ != state_dict["cls_name"]:
             raise ValueError(
-                'Loading a wrong Laplace type. Make sure `subset_of_weights` and'
-                + ' `hessian_structure` are correct!'
+                "Loading a wrong Laplace type. Make sure `subset_of_weights` and"
+                + " `hessian_structure` are correct!"
             )
-        if self.n_params is not None and len(state_dict['mean']) != self.n_params:
+        if self.n_params is not None and len(state_dict["mean"]) != self.n_params:
             raise ValueError(
-                'Attempting to load Laplace with different number of parameters than the model.'
-                + ' Make sure that you use the same `subset_of_weights` value and the same `.requires_grad`'
-                + ' switch on `model.parameters()`.'
+                "Attempting to load Laplace with different number of parameters than the model."
+                + " Make sure that you use the same `subset_of_weights` value and the same `.requires_grad`"
+                + " switch on `model.parameters()`."
             )
-        if self.likelihood != state_dict['likelihood']:
-            raise ValueError('Different likelihoods detected!')
+        if self.likelihood != state_dict["likelihood"]:
+            raise ValueError("Different likelihoods detected!")
 
         # Ignorable warnings
-        if self.prior_mean is None and state_dict['prior_mean'] is not None:
+        if self.prior_mean is None and state_dict["prior_mean"] is not None:
             warnings.warn(
-                'Loading non-`None` prior mean into a `None` prior mean. You might get wrong results.'
+                "Loading non-`None` prior mean into a `None` prior mean. You might get wrong results."
             )
-        if self.temperature != state_dict['temperature']:
+        if self.temperature != state_dict["temperature"]:
             warnings.warn(
-                'Different `temperature` parameters detected. Some calculation might be off!'
+                "Different `temperature` parameters detected. Some calculation might be off!"
             )
-        if self.enable_backprop != state_dict['enable_backprop']:
+        if self.enable_backprop != state_dict["enable_backprop"]:
             warnings.warn(
-                'Different `enable_backprop` values. You might encounter error when differentiating'
-                + ' the predictive mean and variance.'
+                "Different `enable_backprop` values. You might encounter error when differentiating"
+                + " the predictive mean and variance."
             )
 
-        self.mean = state_dict['mean']
-        self.H = state_dict['H']
-        self.loss = state_dict['loss']
-        self.prior_mean = state_dict['prior_mean']
-        self.prior_precision = state_dict['prior_precision']
-        self.sigma_noise = state_dict['sigma_noise']
-        self.n_data = state_dict['n_data']
-        self.n_outputs = state_dict['n_outputs']
-        setattr(self.model, 'output_size', self.n_outputs)
-        self.likelihood = state_dict['likelihood']
-        self.temperature = state_dict['temperature']
-        self.enable_backprop = state_dict['enable_backprop']
+        self.mean = state_dict["mean"]
+        self.H = state_dict["H"]
+        self.loss = state_dict["loss"]
+        self.prior_mean = state_dict["prior_mean"]
+        self.prior_precision = state_dict["prior_precision"]
+        self.sigma_noise = state_dict["sigma_noise"]
+        self.n_data = state_dict["n_data"]
+        self.n_outputs = state_dict["n_outputs"]
+        setattr(self.model, "output_size", self.n_outputs)
+        self.likelihood = state_dict["likelihood"]
+        self.temperature = state_dict["temperature"]
+        self.enable_backprop = state_dict["enable_backprop"]
 
 
 class FullLaplace(ParametricLaplace):
@@ -1127,7 +1137,7 @@ class FullLaplace(ParametricLaplace):
     """
 
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
-    _key = ('all', 'full')
+    _key = ("all", "full")
 
     def __init__(
         self,
@@ -1138,8 +1148,8 @@ def __init__(
         prior_mean=0.0,
         temperature=1.0,
         enable_backprop=False,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
         backend=None,
         backend_kwargs=None,
     ):
@@ -1218,12 +1228,12 @@ def square_norm(self, value):
         return delta @ self.posterior_precision @ delta
 
     def functional_variance(self, Js):
-        return torch.einsum('ncp,pq,nkq->nck', Js, self.posterior_covariance, Js)
+        return torch.einsum("ncp,pq,nkq->nck", Js, self.posterior_covariance, Js)
 
     def functional_covariance(self, Js):
         n_batch, n_outs, n_params = Js.shape
         Js = Js.reshape(n_batch * n_outs, n_params)
-        return torch.einsum('np,pq,mq->nm', Js, self.posterior_covariance, Js)
+        return torch.einsum("np,pq,mq->nm", Js, self.posterior_covariance, Js)
 
     def sample(self, n_samples=100, generator=None):
         samples = torch.randn(
@@ -1248,7 +1258,7 @@ class KronLaplace(ParametricLaplace):
     """
 
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
-    _key = ('all', 'kron')
+    _key = ("all", "kron")
 
     def __init__(
         self,
@@ -1259,8 +1269,8 @@ def __init__(
         prior_mean=0.0,
         temperature=1.0,
         enable_backprop=False,
-        dict_key_x='inputs_id',
-        dict_key_y='labels',
+        dict_key_x="inputs_id",
+        dict_key_y="labels",
         backend=None,
         damping=False,
         backend_kwargs=None,
@@ -1287,8 +1297,8 @@ def _init_H(self):
         self.H = Kron.init_from_model(self.params, self._device)
 
     def _check_H_init(self):
-        if getattr(self, 'H_facs', None) is None:
-            raise AttributeError('Laplace not fitted. Run fit() first.')
+        if getattr(self, "H_facs", None) is None:
+            raise AttributeError("Laplace not fitted. Run fit() first.")
 
     def _curv_closure(self, X, y, N):
         return self.backend.kron(X, y, N=N, **self._asdl_fisher_kwargs)
@@ -1374,18 +1384,18 @@ def prior_precision(self, prior_precision):
         # Extend setter from Laplace to restrict prior precision structure.
         super(KronLaplace, type(self)).prior_precision.fset(self, prior_precision)
         if len(self.prior_precision) not in [1, self.n_layers]:
-            raise ValueError('Prior precision for Kron either scalar or per-layer.')
+            raise ValueError("Prior precision for Kron either scalar or per-layer.")
 
     def state_dict(self) -> dict:
         state_dict = super().state_dict()
-        state_dict['H'] = self.H_facs.kfacs
+        state_dict["H"] = self.H_facs.kfacs
         return state_dict
 
     def load_state_dict(self, state_dict: dict):
         super().load_state_dict(state_dict)
         self._init_H()
         self.H_facs = self.H
-        self.H_facs.kfacs = state_dict['H']
+        self.H_facs.kfacs = state_dict["H"]
         self.H = self.H_facs.decompose(damping=self.damping)
 
 
@@ -1403,7 +1413,7 @@ class LowRankLaplace(ParametricLaplace):
     See `BaseLaplace` for the full interface.
     """
 
-    _key = ('all', 'lowrank')
+    _key = ("all", "lowrank")
 
     def __init__(
         self,
@@ -1414,8 +1424,8 @@ def __init__(
         prior_mean=0,
         temperature=1,
         enable_backprop=False,
-        dict_key_x='inputs_id',
-        dict_key_y='labels',
+        dict_key_x="inputs_id",
+        dict_key_y="labels",
         backend=AsdfghjklHessian,
         backend_kwargs=None,
     ):
@@ -1438,19 +1448,19 @@ def _init_H(self):
 
     @property
     def V(self):
-        (U, l), prior_prec_diag = self.posterior_precision
+        (U, _), prior_prec_diag = self.posterior_precision
         return U / prior_prec_diag.reshape(-1, 1)
 
     @property
     def Kinv(self):
-        (U, l), _ = self.posterior_precision
-        return torch.inverse(torch.diag(1 / l) + U.T @ self.V)
+        (U, eigval), _ = self.posterior_precision
+        return torch.inverse(torch.diag(1 / eigval) + U.T @ self.V)
 
     def fit(self, train_loader, override=True):
         # override fit since output of eighessian not additive across batch
         if not override:
             # LowRankLA cannot be updated since eigenvalue representation not additive
-            raise ValueError('LowRank LA does not support updating.')
+            raise ValueError("LowRank LA does not support updating.")
 
         self.model.eval()
         self.mean = parameters_to_vector(self.model.parameters())
@@ -1465,7 +1475,7 @@ def fit(self, train_loader, override=True):
             except (TypeError, AttributeError):
                 out = self.model(X.to(self._device))
         self.n_outputs = out.shape[-1]
-        setattr(self.model, 'output_size', self.n_outputs)
+        setattr(self.model, "output_size", self.n_outputs)
 
         eigenvectors, eigenvalues, loss = self.backend.eig_lowrank(train_loader)
         self.H = (eigenvectors, eigenvalues)
@@ -1489,17 +1499,17 @@ def posterior_precision(self):
         return (self.H[0], self._H_factor * self.H[1]), self.prior_precision_diag
 
     def functional_variance(self, Jacs):
-        prior_var = torch.einsum('ncp,nkp->nck', Jacs / self.prior_precision_diag, Jacs)
-        Jacs_V = torch.einsum('ncp,pl->ncl', Jacs, self.V)
-        info_gain = torch.einsum('ncl,nkl->nck', Jacs_V @ self.Kinv, Jacs_V)
+        prior_var = torch.einsum("ncp,nkp->nck", Jacs / self.prior_precision_diag, Jacs)
+        Jacs_V = torch.einsum("ncp,pl->ncl", Jacs, self.V)
+        info_gain = torch.einsum("ncl,nkl->nck", Jacs_V @ self.Kinv, Jacs_V)
         return prior_var - info_gain
 
     def functional_covariance(self, Jacs):
         n_batch, n_outs, n_params = Jacs.shape
         Jacs = Jacs.reshape(n_batch * n_outs, n_params)
-        prior_cov = torch.einsum('np,mp->nm', Jacs / self.prior_precision_diag, Jacs)
-        Jacs_V = torch.einsum('np,pl->nl', Jacs, self.V)
-        info_gain = torch.einsum('nl,ml->nm', Jacs_V @ self.Kinv, Jacs_V)
+        prior_cov = torch.einsum("np,mp->nm", Jacs / self.prior_precision_diag, Jacs)
+        Jacs_V = torch.einsum("np,pl->nl", Jacs, self.V)
+        info_gain = torch.einsum("nl,ml->nm", Jacs_V @ self.Kinv, Jacs_V)
         cov = prior_cov - info_gain
         assert cov.shape == (n_batch * n_outs, n_batch * n_outs)
         return cov
@@ -1522,8 +1532,10 @@ def sample(self, n_samples, generator=None):
 
     @property
     def log_det_posterior_precision(self):
-        (U, l), prior_prec_diag = self.posterior_precision
-        return l.log().sum() + prior_prec_diag.log().sum() - torch.logdet(self.Kinv)
+        (U, eigval), prior_prec_diag = self.posterior_precision
+        return (
+            eigval.log().sum() + prior_prec_diag.log().sum() - torch.logdet(self.Kinv)
+        )
 
 
 class DiagLaplace(ParametricLaplace):
@@ -1534,7 +1546,7 @@ class DiagLaplace(ParametricLaplace):
     """
 
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
-    _key = ('all', 'diag')
+    _key = ("all", "diag")
 
     def _init_H(self):
         self.H = torch.zeros(self.n_params, device=self._device)
@@ -1586,13 +1598,13 @@ def square_norm(self, value):
 
     def functional_variance(self, Js: torch.Tensor) -> torch.Tensor:
         self._check_jacobians(Js)
-        return torch.einsum('ncp,p,nkp->nck', Js, self.posterior_variance, Js)
+        return torch.einsum("ncp,p,nkp->nck", Js, self.posterior_variance, Js)
 
     def functional_covariance(self, Js):
         self._check_jacobians(Js)
         n_batch, n_outs, n_params = Js.shape
         Js = Js.reshape(n_batch * n_outs, n_params)
-        cov = torch.einsum('np,p,mp->nm', Js, self.posterior_variance, Js)
+        cov = torch.einsum("np,p,mp->nm", Js, self.posterior_variance, Js)
         return cov
 
     def sample(self, n_samples=100, generator=None):
diff --git a/laplace/curvature/__init__.py b/laplace/curvature/__init__.py
index 7bc83ab..ed7f054 100644
--- a/laplace/curvature/__init__.py
+++ b/laplace/curvature/__init__.py
@@ -1,54 +1,54 @@
 import logging
 
-from laplace.curvature.curvature import CurvatureInterface, GGNInterface, EFInterface
+from laplace.curvature.curvature import CurvatureInterface, EFInterface, GGNInterface
 
 try:
-    from laplace.curvature.backpack import BackPackGGN, BackPackEF, BackPackInterface
+    from laplace.curvature.backpack import BackPackEF, BackPackGGN, BackPackInterface
 except ModuleNotFoundError:
-    logging.info('Backpack backend not available.')
+    logging.info("Backpack backend not available.")
 
 try:
     from laplace.curvature.asdfghjkl import (
-        AsdfghjklHessian,
-        AsdfghjklGGN,
         AsdfghjklEF,
+        AsdfghjklGGN,
+        AsdfghjklHessian,
         AsdfghjklInterface,
     )
 except ModuleNotFoundError:
-    logging.info('Asdfghjkl backend not available.')
+    logging.info("Asdfghjkl backend not available.")
 
 try:
-    from laplace.curvature.asdl import AsdlHessian, AsdlGGN, AsdlEF, AsdlInterface
+    from laplace.curvature.asdl import AsdlEF, AsdlGGN, AsdlHessian, AsdlInterface
 except ModuleNotFoundError:
-    logging.info('ASDL backend not available.')
+    logging.info("ASDL backend not available.")
 
 try:
     from laplace.curvature.curvlinops import (
-        CurvlinopsHessian,
-        CurvlinopsGGN,
         CurvlinopsEF,
+        CurvlinopsGGN,
+        CurvlinopsHessian,
         CurvlinopsInterface,
     )
 except ModuleNotFoundError:
-    logging.info('Curvlinops backend not available.')
+    logging.info("Curvlinops backend not available.")
 
 __all__ = [
-    'CurvatureInterface',
-    'GGNInterface',
-    'EFInterface',
-    'BackPackInterface',
-    'BackPackGGN',
-    'BackPackEF',
-    'AsdfghjklInterface',
-    'AsdfghjklGGN',
-    'AsdfghjklEF',
-    'AsdfghjklHessian',
-    'AsdlInterface',
-    'AsdlGGN',
-    'AsdlEF',
-    'AsdlHessian',
-    'CurvlinopsInterface',
-    'CurvlinopsGGN',
-    'CurvlinopsEF',
-    'CurvlinopsHessian',
+    "CurvatureInterface",
+    "GGNInterface",
+    "EFInterface",
+    "BackPackInterface",
+    "BackPackGGN",
+    "BackPackEF",
+    "AsdfghjklInterface",
+    "AsdfghjklGGN",
+    "AsdfghjklEF",
+    "AsdfghjklHessian",
+    "AsdlInterface",
+    "AsdlGGN",
+    "AsdlEF",
+    "AsdlHessian",
+    "CurvlinopsInterface",
+    "CurvlinopsGGN",
+    "CurvlinopsEF",
+    "CurvlinopsHessian",
 ]
diff --git a/laplace/curvature/asdfghjkl.py b/laplace/curvature/asdfghjkl.py
index 9fff7a8..b83009d 100644
--- a/laplace/curvature/asdfghjkl.py
+++ b/laplace/curvature/asdfghjkl.py
@@ -1,14 +1,20 @@
 import warnings
+
 import numpy as np
 import torch
-
-from asdfghjkl import FISHER_EXACT, FISHER_MC, COV
-from asdfghjkl import SHAPE_KRON, SHAPE_DIAG, SHAPE_FULL
-from asdfghjkl import fisher_for_cross_entropy
-from asdfghjkl.hessian import hessian_eigenvalues, hessian_for_loss
+from asdfghjkl import (
+    COV,
+    FISHER_EXACT,
+    FISHER_MC,
+    SHAPE_DIAG,
+    SHAPE_FULL,
+    SHAPE_KRON,
+    fisher_for_cross_entropy,
+)
 from asdfghjkl.gradient import batch_gradient
+from asdfghjkl.hessian import hessian_eigenvalues, hessian_for_loss
 
-from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
+from laplace.curvature import CurvatureInterface, EFInterface, GGNInterface
 from laplace.utils import Kron, _is_batchnorm
 
 EPS = 1e-6
@@ -80,24 +86,24 @@ def _get_kron_factors(self, curv, M):
         kfacs = list()
         for module in curv._model.modules():
             if _is_batchnorm(module):
-                warnings.warn('BatchNorm unsupported for Kron, ignore.')
+                warnings.warn("BatchNorm unsupported for Kron, ignore.")
                 continue
 
             stats = getattr(module, self._ggn_type, None)
             if stats is None:
                 continue
-            if hasattr(module, 'bias') and module.bias is not None:
+            if hasattr(module, "bias") and module.bias is not None:
                 # split up bias and weights
                 kfacs.append([stats.kron.B, stats.kron.A[:-1, :-1]])
                 kfacs.append([stats.kron.B * stats.kron.A[-1, -1] / M])
-            elif hasattr(module, 'weight'):
+            elif hasattr(module, "weight"):
                 p, q = stats.kron.B.numel(), stats.kron.A.numel()
                 if p == q == 1:
                     kfacs.append([stats.kron.B * stats.kron.A])
                 else:
                     kfacs.append([stats.kron.B, stats.kron.A])
             else:
-                raise ValueError(f'Whats happening with {module}?')
+                raise ValueError(f"Whats happening with {module}?")
         return Kron(kfacs)
 
     @staticmethod
@@ -144,8 +150,8 @@ def __init__(
         model,
         likelihood,
         last_layer=False,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
         low_rank=10,
     ):
         super().__init__(
@@ -153,8 +159,8 @@ def __init__(
             likelihood,
             last_layer,
             None,
-            dict_key_x='input_ids',
-            dict_key_y='labels',
+            dict_key_x="input_ids",
+            dict_key_y="labels",
         )
         self.low_rank = low_rank
 
@@ -202,12 +208,12 @@ def __init__(
         likelihood,
         last_layer=False,
         subnetwork_indices=None,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
         stochastic=False,
     ):
-        if likelihood != 'classification':
-            raise ValueError('This backend only supports classification currently.')
+        if likelihood != "classification":
+            raise ValueError("This backend only supports classification currently.")
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
         )
@@ -226,11 +232,11 @@ def __init__(
         model,
         likelihood,
         last_layer=False,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
     ):
-        if likelihood != 'classification':
-            raise ValueError('This backend only supports classification currently.')
+        if likelihood != "classification":
+            raise ValueError("This backend only supports classification currently.")
         super().__init__(model, likelihood, last_layer, None, dict_key_x, dict_key_y)
 
     @property
@@ -248,12 +254,12 @@ def _flatten_after_batch(tensor: torch.Tensor):
 def _get_batch_grad(model):
     batch_grads = list()
     for module in model.modules():
-        if hasattr(module, 'op_results'):
-            res = module.op_results['batch_grads']
-            if 'weight' in res:
-                batch_grads.append(_flatten_after_batch(res['weight']))
-            if 'bias' in res:
-                batch_grads.append(_flatten_after_batch(res['bias']))
-            if len(set(res.keys()) - {'weight', 'bias'}) > 0:
-                raise ValueError(f'Invalid parameter keys {res.keys()}')
+        if hasattr(module, "op_results"):
+            res = module.op_results["batch_grads"]
+            if "weight" in res:
+                batch_grads.append(_flatten_after_batch(res["weight"]))
+            if "bias" in res:
+                batch_grads.append(_flatten_after_batch(res["bias"]))
+            if len(set(res.keys()) - {"weight", "bias"}) > 0:
+                raise ValueError(f"Invalid parameter keys {res.keys()}")
     return torch.cat(batch_grads, dim=1)
diff --git a/laplace/curvature/asdl.py b/laplace/curvature/asdl.py
index b62fa4f..8765855 100644
--- a/laplace/curvature/asdl.py
+++ b/laplace/curvature/asdl.py
@@ -1,23 +1,21 @@
-from collections.abc import MutableMapping
 import warnings
+from collections.abc import MutableMapping
 
-import numpy as np
 import torch
-
+from asdl.fisher import FisherConfig, get_fisher_maker
+from asdl.grad_maker import LOSS_CROSS_ENTROPY, LOSS_MSE
+from asdl.gradient import batch_gradient
+from asdl.hessian import HessianConfig, HessianMaker
 from asdl.matrices import (
+    FISHER_EMP,
     FISHER_EXACT,
     FISHER_MC,
-    FISHER_EMP,
-    SHAPE_KRON,
     SHAPE_DIAG,
     SHAPE_FULL,
+    SHAPE_KRON,
 )
-from asdl.grad_maker import LOSS_MSE, LOSS_CROSS_ENTROPY
-from asdl.fisher import FisherConfig, get_fisher_maker
-from asdl.hessian import HessianMaker, HessianConfig
-from asdl.gradient import batch_gradient
 
-from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
+from laplace.curvature import CurvatureInterface, EFInterface, GGNInterface
 from laplace.utils import Kron, _is_batchnorm
 
 EPS = 1e-6
@@ -32,8 +30,8 @@ def __init__(
         likelihood,
         last_layer=False,
         subnetwork_indices=None,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
     ):
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
@@ -41,7 +39,7 @@ def __init__(
 
     @property
     def loss_type(self):
-        return LOSS_MSE if self.likelihood == 'regression' else LOSS_CROSS_ENTROPY
+        return LOSS_MSE if self.likelihood == "regression" else LOSS_CROSS_ENTROPY
 
     def jacobians(self, x, enable_backprop=False):
         """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\)
@@ -125,25 +123,25 @@ def _get_kron_factors(self, M):
         kfacs = list()
         for module in self.model.modules():
             if _is_batchnorm(module):
-                warnings.warn('BatchNorm unsupported for Kron, ignore.')
+                warnings.warn("BatchNorm unsupported for Kron, ignore.")
                 continue
 
-            stats = getattr(module, 'fisher', None)
+            stats = getattr(module, "fisher", None)
             if stats is None:
                 continue
 
-            if hasattr(module, 'bias') and module.bias is not None:
+            if hasattr(module, "bias") and module.bias is not None:
                 # split up bias and weights
                 kfacs.append([stats.kron.B, stats.kron.A])
                 kfacs.append([stats.kron.B])
-            elif hasattr(module, 'weight'):
+            elif hasattr(module, "weight"):
                 p, q = stats.kron.B.numel(), stats.kron.A.numel()
                 if p == q == 1:
                     kfacs.append([stats.kron.B * stats.kron.A])
                 else:
                     kfacs.append([stats.kron.B, stats.kron.A])
             else:
-                raise ValueError(f'Whats happening with {module}?')
+                raise ValueError(f"Whats happening with {module}?")
         return Kron(kfacs)
 
     @staticmethod
@@ -166,7 +164,7 @@ def diag(self, X, y, N=None, **kwargs):
         )
         fisher_maker = get_fisher_maker(self.model, cfg)
         y = y if self.loss_type == LOSS_MSE else y.view(-1)
-        if 'emp' in self._ggn_type:
+        if "emp" in self._ggn_type:
             dummy = fisher_maker.setup_model_call(self._model, X)
             dummy = (
                 dummy if self.loss_type == LOSS_MSE else dummy.view(-1, dummy.size(-1))
@@ -180,14 +178,14 @@ def diag(self, X, y, N=None, **kwargs):
         loss = self.lossfunc(f.detach(), y)
         vec = list()
         for module in self.model.modules():
-            stats = getattr(module, 'fisher', None)
+            stats = getattr(module, "fisher", None)
             if stats is None:
                 continue
             vec.extend(stats.to_vector())
         diag_ggn = torch.cat(vec)
         if self.subnetwork_indices is not None:
             diag_ggn = diag_ggn[self.subnetwork_indices]
-        if type(self) is AsdlEF and self.likelihood == 'regression':
+        if type(self) is AsdlEF and self.likelihood == "regression":
             curv_factor = 0.5  # correct scaling for diag ef
         else:
             curv_factor = 1.0  # ASDL uses proper 1/2 * MSELoss
@@ -205,7 +203,7 @@ def kron(self, X, y, N, **kwargs):
         )
         fisher_maker = get_fisher_maker(self.model, cfg)
         y = y if self.loss_type == LOSS_MSE else y.view(-1)
-        if 'emp' in self._ggn_type:
+        if "emp" in self._ggn_type:
             dummy = fisher_maker.setup_model_call(self._model, X)
             dummy = (
                 dummy if self.loss_type == LOSS_MSE else dummy.view(-1, dummy.size(-1))
@@ -220,7 +218,7 @@ def kron(self, X, y, N, **kwargs):
         M = len(y)
         kron = self._get_kron_factors(M)
         kron = self._rescale_kron_factors(kron, N)
-        if type(self) is AsdlEF and self.likelihood == 'regression':
+        if type(self) is AsdlEF and self.likelihood == "regression":
             curv_factor = 0.5  # correct scaling for diag ef
         else:
             curv_factor = 1.0  # ASDL uses proper 1/2 * MSELoss
@@ -244,8 +242,8 @@ def __init__(
         model,
         likelihood,
         last_layer=False,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
     ):
         super().__init__(
             model,
@@ -287,8 +285,8 @@ def __init__(
         likelihood,
         last_layer=False,
         subnetwork_indices=None,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
         stochastic=False,
     ):
         super().__init__(
@@ -309,8 +307,8 @@ def __init__(
         model,
         likelihood,
         last_layer=False,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
     ):
         super().__init__(model, likelihood, last_layer, None, dict_key_x, dict_key_y)
 
diff --git a/laplace/curvature/backpack.py b/laplace/curvature/backpack.py
index 6b66015..b6c0195 100644
--- a/laplace/curvature/backpack.py
+++ b/laplace/curvature/backpack.py
@@ -1,18 +1,18 @@
 from typing import Tuple
-import torch
 
+import torch
 from backpack import backpack, extend, memory_cleanup
+from backpack.context import CTX
 from backpack.extensions import (
-    DiagGGNExact,
-    DiagGGNMC,
     KFAC,
     KFLR,
-    SumGradSquared,
     BatchGrad,
+    DiagGGNExact,
+    DiagGGNMC,
+    SumGradSquared,
 )
-from backpack.context import CTX
 
-from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
+from laplace.curvature import CurvatureInterface, EFInterface, GGNInterface
 from laplace.utils import Kron
 
 
@@ -25,8 +25,8 @@ def __init__(
         likelihood,
         last_layer=False,
         subnetwork_indices=None,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
     ):
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
@@ -71,7 +71,7 @@ def jacobians(self, x, enable_backprop=False):
                 to_cat = []
                 for param in model.parameters():
                     to_cat.append(param.grad_batch.reshape(x.shape[0], -1))
-                    delattr(param, 'grad_batch')
+                    delattr(param, "grad_batch")
                 Jk = torch.cat(to_cat, dim=1)
                 if self.subnetwork_indices is not None:
                     Jk = Jk[:, self.subnetwork_indices]
@@ -126,8 +126,8 @@ def __init__(
         likelihood,
         last_layer=False,
         subnetwork_indices=None,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
         stochastic=False,
     ):
         super().__init__(
@@ -164,8 +164,8 @@ def diag(self, X, y, **kwargs):
         context = DiagGGNMC if self.stochastic else DiagGGNExact
         f = self.model(X)
         # Assumes that the last dimension of f is of size outputs.
-        f = f if self.likelihood == 'regression' else f.view(-1, f.size(-1))
-        y = y if self.likelihood == 'regression' else y.view(-1)
+        f = f if self.likelihood == "regression" else f.view(-1, f.size(-1))
+        y = y if self.likelihood == "regression" else y.view(-1)
         loss = self.lossfunc(f, y)
         with backpack(context()):
             loss.backward()
@@ -179,8 +179,8 @@ def kron(self, X, y, N, **kwargs) -> Tuple[torch.Tensor, Kron]:
         context = KFAC if self.stochastic else KFLR
         f = self.model(X)
         # Assumes that the last dimension of f is of size outputs.
-        f = f if self.likelihood == 'regression' else f.view(-1, f.size(-1))
-        y = y if self.likelihood == 'regression' else y.view(-1)
+        f = f if self.likelihood == "regression" else f.view(-1, f.size(-1))
+        y = y if self.likelihood == "regression" else y.view(-1)
         loss = self.lossfunc(f, y)
         with backpack(context()):
             loss.backward()
@@ -196,8 +196,8 @@ class BackPackEF(BackPackInterface, EFInterface):
     def diag(self, X, y, **kwargs):
         f = self.model(X)
         # Assumes that the last dimension of f is of size outputs.
-        f = f if self.likelihood == 'regression' else f.view(-1, f.size(-1))
-        y = y if self.likelihood == 'regression' else y.view(-1)
+        f = f if self.likelihood == "regression" else f.view(-1, f.size(-1))
+        y = y if self.likelihood == "regression" else y.view(-1)
         loss = self.lossfunc(f, y)
         with backpack(SumGradSquared()):
             loss.backward()
@@ -210,12 +210,12 @@ def diag(self, X, y, **kwargs):
         return self.factor * loss.detach(), self.factor * diag_EF
 
     def kron(self, X, y, **kwargs):
-        raise NotImplementedError('Unavailable through Backpack.')
+        raise NotImplementedError("Unavailable through Backpack.")
 
 
 def _cleanup(module):
     for child in module.children():
         _cleanup(child)
 
-    setattr(module, '_backpack_extend', False)
+    setattr(module, "_backpack_extend", False)
     memory_cleanup(module)
diff --git a/laplace/curvature/curvature.py b/laplace/curvature/curvature.py
index 6d1e040..779bf0f 100644
--- a/laplace/curvature/curvature.py
+++ b/laplace/curvature/curvature.py
@@ -1,6 +1,5 @@
 import torch
-from torch.nn import MSELoss, CrossEntropyLoss
-from torch.nn.utils import parameters_to_vector, vector_to_parameters
+from torch.nn import CrossEntropyLoss, MSELoss
 
 
 class CurvatureInterface:
@@ -43,10 +42,10 @@ def __init__(
         likelihood,
         last_layer=False,
         subnetwork_indices=None,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
     ):
-        assert likelihood in ['regression', 'classification']
+        assert likelihood in ["regression", "classification"]
         self.likelihood = likelihood
         self.model = model
         self.last_layer = last_layer
@@ -54,11 +53,11 @@ def __init__(
         self.dict_key_x = dict_key_x
         self.dict_key_y = dict_key_y
 
-        if likelihood == 'regression':
-            self.lossfunc = MSELoss(reduction='sum')
+        if likelihood == "regression":
+            self.lossfunc = MSELoss(reduction="sum")
             self.factor = 0.5
         else:
-            self.lossfunc = CrossEntropyLoss(reduction='sum')
+            self.lossfunc = CrossEntropyLoss(reduction="sum")
             self.factor = 1.0
 
         self.params = [p for p in self._model.parameters() if p.requires_grad]
@@ -177,7 +176,7 @@ def last_layer_jacobians(self, x, enable_backprop=False):
             .tile(bsize, 1, 1)
         )
         # Jacobians are batch x output x params
-        Js = torch.einsum('kp,kij->kijp', phi, identity).reshape(bsize, output_size, -1)
+        Js = torch.einsum("kp,kij->kijp", phi, identity).reshape(bsize, output_size, -1)
         if self.model.last_layer.bias is not None:
             Js = torch.cat([Js, identity], dim=2)
 
@@ -321,8 +320,8 @@ def __init__(
         likelihood,
         last_layer=False,
         subnetwork_indices=None,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
         stochastic=False,
         num_samples=1,
     ):
@@ -339,7 +338,7 @@ def _get_mc_functional_fisher(self, f):
         F = 0
 
         for _ in range(self.num_samples):
-            if self.likelihood == 'regression':
+            if self.likelihood == "regression":
                 y_sample = f + torch.randn(f.shape, device=f.device)  # N(y | f, 1)
                 grad_sample = f - y_sample  # functional MSE grad
             else:  # classification with softmax
@@ -351,18 +350,18 @@ def _get_mc_functional_fisher(self, f):
             F += (
                 1
                 / self.num_samples
-                * torch.einsum('bc,bk->bck', grad_sample, grad_sample)
+                * torch.einsum("bc,bk->bck", grad_sample, grad_sample)
             )
 
         return F
 
     def _get_functional_hessian(self, f):
-        if self.likelihood == 'regression':
+        if self.likelihood == "regression":
             return None
         else:
             # second derivative of log lik is diag(p) - pp^T
             ps = torch.softmax(f, dim=-1)
-            G = torch.diag_embed(ps) - torch.einsum('mk,mc->mck', ps, ps)
+            G = torch.diag_embed(ps) - torch.einsum("mk,mc->mck", ps, ps)
             return G
 
     def full(self, x, y, **kwargs):
@@ -391,9 +390,9 @@ def full(self, x, y, **kwargs):
         )
 
         if H_lik is not None:
-            H = torch.einsum('bcp,bck,bkq->pq', Js, H_lik, Js)
+            H = torch.einsum("bcp,bck,bkq->pq", Js, H_lik, Js)
         else:  # The case of exact GGN for regression
-            H = torch.einsum('bcp,bcq->pq', Js, Js)
+            H = torch.einsum("bcp,bcq->pq", Js, Js)
         loss = self.factor * self.lossfunc(f, y)
 
         return loss.detach(), H.detach()
@@ -409,9 +408,9 @@ def diag(self, X, y, **kwargs):
         )
 
         if H_lik is not None:
-            H = torch.einsum('bcp,bck,bkp->p', Js, H_lik, Js)
+            H = torch.einsum("bcp,bck,bkp->p", Js, H_lik, Js)
         else:  # The case of exact GGN for regression
-            H = torch.einsum('bcp,bcp->p', Js, Js)
+            H = torch.einsum("bcp,bcp->p", Js, Js)
 
         return loss.detach(), H.detach()
 
@@ -466,11 +465,11 @@ def full(self, x, y, **kwargs):
             EF `(parameters, parameters)`
         """
         Gs, loss = self.gradients(x, y)
-        H_ef = torch.einsum('bp,bq->pq', Gs, Gs)
+        H_ef = torch.einsum("bp,bq->pq", Gs, Gs)
         return self.factor * loss.detach(), self.factor * H_ef
 
     def diag(self, X, y, **kwargs):
         # Gs is (batchsize, n_params)
         Gs, loss = self.gradients(X, y)
-        diag_ef = torch.einsum('bp,bp->p', Gs, Gs)
+        diag_ef = torch.einsum("bp,bp->p", Gs, Gs)
         return self.factor * loss.detach(), self.factor * diag_ef
diff --git a/laplace/curvature/curvlinops.py b/laplace/curvature/curvlinops.py
index d7bb9d7..d3930ae 100644
--- a/laplace/curvature/curvlinops.py
+++ b/laplace/curvature/curvlinops.py
@@ -1,19 +1,17 @@
-import torch
-import numpy as np
+from collections import UserDict
 
+import torch
 from curvlinops import (
-    HessianLinearOperator,
-    GGNLinearOperator,
-    FisherMCLinearOperator,
     EFLinearOperator,
+    FisherMCLinearOperator,
+    GGNLinearOperator,
+    HessianLinearOperator,
     KFACLinearOperator,
 )
 
-from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
+from laplace.curvature import CurvatureInterface, EFInterface, GGNInterface
 from laplace.utils import Kron
 
-from collections import UserDict
-
 
 class CurvlinopsInterface(CurvatureInterface):
     """Interface for Curvlinops backend. <https://github.com/f-dangel/curvlinops>"""
@@ -24,8 +22,8 @@ def __init__(
         likelihood,
         last_layer=False,
         subnetwork_indices=None,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
     ):
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
@@ -57,22 +55,23 @@ def _get_kron_factors(self, linop):
             A = linop._input_covariances[name]
             B = linop._gradient_covariances[name]
 
-            if hasattr(module, 'bias') and module.bias is not None:
+            if hasattr(module, "bias") and module.bias is not None:
                 kfacs.append([B, A])
                 kfacs.append([B])
-            elif hasattr(module, 'weight'):
+            elif hasattr(module, "weight"):
                 p, q = B.numel(), A.numel()
                 if p == q == 1:
                     kfacs.append([B * A])
                 else:
                     kfacs.append([B, A])
             else:
-                raise ValueError(f'Whats happening with {module}?')
+                raise ValueError(f"Whats happening with {module}?")
         return Kron(kfacs)
 
     def kron(self, X, y, N, **kwargs):
         if isinstance(X, (dict, UserDict)):
-            kwargs['batch_size_fn'] = lambda x: x[self.dict_key_x].shape[0]
+            kwargs["batch_size_fn"] = lambda x: x[self.dict_key_x].shape[0]
+
         linop = KFACLinearOperator(
             self.model,
             self.lossfunc,
@@ -102,9 +101,9 @@ def full(self, X, y, **kwargs):
         if self.subnetwork_indices is not None:
             return super().full(X, y, **kwargs)
 
-        curvlinops_kwargs = {k: v for k, v in kwargs.items() if k != 'N'}
+        curvlinops_kwargs = {k: v for k, v in kwargs.items() if k != "N"}
         if isinstance(X, (dict, UserDict)):
-            curvlinops_kwargs['batch_size_fn'] = lambda x: x[self.dict_key_x].shape[0]
+            curvlinops_kwargs["batch_size_fn"] = lambda x: x[self.dict_key_x].shape[0]
 
         linop = self._linop_context(
             self.model,
@@ -134,8 +133,8 @@ def __init__(
         likelihood,
         last_layer=False,
         subnetwork_indices=None,
-        dict_key_x='input_ids',
-        dict_key_y='labels',
+        dict_key_x="input_ids",
+        dict_key_y="labels",
         stochastic=False,
     ):
         super().__init__(
@@ -145,7 +144,7 @@ def __init__(
 
     @property
     def _kron_fisher_type(self):
-        return 'mc' if self.stochastic else 'type-2'
+        return "mc" if self.stochastic else "type-2"
 
     @property
     def _linop_context(self):
@@ -157,7 +156,7 @@ class CurvlinopsEF(CurvlinopsInterface, EFInterface):
 
     @property
     def _kron_fisher_type(self):
-        return 'empirical'
+        return "empirical"
 
     @property
     def _linop_context(self):
diff --git a/laplace/laplace.py b/laplace/laplace.py
index f8c4577..6a0d598 100644
--- a/laplace/laplace.py
+++ b/laplace/laplace.py
@@ -1,12 +1,12 @@
+from laplace import *  # noqa F403
 from laplace.baselaplace import ParametricLaplace
-from laplace import *
 
 
 def Laplace(
     model,
     likelihood,
-    subset_of_weights='last_layer',
-    hessian_structure='kron',
+    subset_of_weights="last_layer",
+    hessian_structure="kron",
     *args,
     **kwargs,
 ):
@@ -26,15 +26,15 @@ def Laplace(
     laplace : ParametricLaplace
         chosen subclass of ParametricLaplace instantiated with additional arguments
     """
-    if subset_of_weights == 'subnetwork' and hessian_structure not in ['full', 'diag']:
+    if subset_of_weights == "subnetwork" and hessian_structure not in ["full", "diag"]:
         raise ValueError(
-            'Subnetwork Laplace requires a full or diagonal Hessian approximation!'
+            "Subnetwork Laplace requires a full or diagonal Hessian approximation!"
         )
 
     laplace_map = {
         subclass._key: subclass
         for subclass in _all_subclasses(ParametricLaplace)
-        if hasattr(subclass, '_key')
+        if hasattr(subclass, "_key")
     }
     laplace_class = laplace_map[(subset_of_weights, hessian_structure)]
     return laplace_class(model, likelihood, *args, **kwargs)
diff --git a/laplace/lllaplace.py b/laplace/lllaplace.py
index fedd010..2bd87c8 100644
--- a/laplace/lllaplace.py
+++ b/laplace/lllaplace.py
@@ -5,12 +5,10 @@
 import torch
 from torch.nn.utils import parameters_to_vector, vector_to_parameters
 
-from laplace.baselaplace import (DiagLaplace, FullLaplace, KronLaplace,
-                                 ParametricLaplace)
+from laplace.baselaplace import DiagLaplace, FullLaplace, KronLaplace, ParametricLaplace
 from laplace.utils import FeatureExtractor, Kron
-from laplace.utils.feature_extractor import FeatureReduction
 
-__all__ = ['LLLaplace', 'FullLLLaplace', 'KronLLLaplace', 'DiagLLLaplace']
+__all__ = ["LLLaplace", "FullLLLaplace", "KronLLLaplace", "DiagLLLaplace"]
 
 
 class LLLaplace(ParametricLaplace):
@@ -89,15 +87,15 @@ def __init__(
         temperature=1.0,
         enable_backprop=False,
         feature_reduction=None,
-        dict_key_x='inputs_id',
-        dict_key_y='labels',
+        dict_key_x="inputs_id",
+        dict_key_y="labels",
         backend=None,
         last_layer_name=None,
         backend_kwargs=None,
         asdl_fisher_kwargs=None,
     ):
         if asdl_fisher_kwargs is not None:
-            raise ValueError('Last-layer Laplace does not support asdl_fisher_kwargs.')
+            raise ValueError("Last-layer Laplace does not support asdl_fisher_kwargs.")
 
         self.H = None
         super().__init__(
@@ -135,7 +133,7 @@ def __init__(
             self.prior_mean = prior_mean
             self.mean = self.prior_mean
             self._init_H()
-        self._backend_kwargs['last_layer'] = True
+        self._backend_kwargs["last_layer"] = True
         self._last_layer_name = last_layer_name
 
     def fit(self, train_loader, override=True):
@@ -152,7 +150,7 @@ def fit(self, train_loader, override=True):
         """
         if not override:
             raise ValueError(
-                'Last-layer Laplace approximations do not support `override=False`.'
+                "Last-layer Laplace approximations do not support `override=False`."
             )
 
         self.model.eval()
@@ -210,7 +208,7 @@ def _nn_predictive_samples(self, X, n_samples=100, generator=None, **model_kwarg
         vector_to_parameters(self.mean, self.model.last_layer.parameters())
         fs = torch.stack(fs)
 
-        if self.likelihood == 'classification':
+        if self.likelihood == "classification":
             fs = torch.softmax(fs, dim=-1)
 
         return fs
@@ -254,19 +252,19 @@ def prior_precision_diag(self):
             return self.prior_precision
 
         else:
-            raise ValueError('Mismatch of prior and model. Diagonal or scalar prior.')
+            raise ValueError("Mismatch of prior and model. Diagonal or scalar prior.")
 
     def state_dict(self) -> dict:
         state_dict = super().state_dict()
-        state_dict['data'] = getattr(self, 'data', None)  # None if not present
-        state_dict['_last_layer_name'] = self._last_layer_name
+        state_dict["data"] = getattr(self, "data", None)  # None if not present
+        state_dict["_last_layer_name"] = self._last_layer_name
         return state_dict
 
     def load_state_dict(self, state_dict: dict):
-        if self._last_layer_name != state_dict['_last_layer_name']:
-            raise ValueError('Different `last_layer_name` detected!')
+        if self._last_layer_name != state_dict["_last_layer_name"]:
+            raise ValueError("Different `last_layer_name` detected!")
 
-        self.data = state_dict['data']
+        self.data = state_dict["data"]
         if self.data is not None:
             self._find_last_layer(self.data)
 
@@ -298,7 +296,7 @@ class FullLLLaplace(LLLaplace, FullLaplace):
     """
 
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
-    _key = ('last_layer', 'full')
+    _key = ("last_layer", "full")
 
 
 class KronLLLaplace(LLLaplace, KronLaplace):
@@ -315,7 +313,7 @@ class KronLLLaplace(LLLaplace, KronLaplace):
     """
 
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
-    _key = ('last_layer', 'kron')
+    _key = ("last_layer", "kron")
 
     def __init__(
         self,
@@ -327,8 +325,8 @@ def __init__(
         temperature=1.0,
         enable_backprop=False,
         feature_reduction=None,
-        dict_key_x='inputs_id',
-        dict_key_y='labels',
+        dict_key_x="inputs_id",
+        dict_key_y="labels",
         backend=None,
         last_layer_name=None,
         damping=False,
@@ -363,4 +361,4 @@ class DiagLLLaplace(LLLaplace, DiagLaplace):
     """
 
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
-    _key = ('last_layer', 'diag')
+    _key = ("last_layer", "diag")
diff --git a/laplace/marglik_training.py b/laplace/marglik_training.py
index c968823..9e8e2df 100644
--- a/laplace/marglik_training.py
+++ b/laplace/marglik_training.py
@@ -1,13 +1,14 @@
+import logging
+import warnings
 from collections.abc import MutableMapping
 from copy import deepcopy
+
 import numpy as np
 import torch
-from torch.optim import Adam
+import tqdm
 from torch.nn import CrossEntropyLoss, MSELoss
 from torch.nn.utils import parameters_to_vector
-import warnings
-import logging
-import tqdm
+from torch.optim import Adam
 
 from laplace import Laplace
 from laplace.curvature import AsdlGGN
@@ -17,8 +18,8 @@
 def marglik_training(
     model,
     train_loader,
-    likelihood='classification',
-    hessian_structure='kron',
+    likelihood="classification",
+    hessian_structure="kron",
     backend=AsdlGGN,
     optimizer_cls=Adam,
     optimizer_kwargs=None,
@@ -26,7 +27,7 @@ def marglik_training(
     scheduler_kwargs=None,
     n_epochs=300,
     lr_hyp=1e-1,
-    prior_structure='layerwise',
+    prior_structure="layerwise",
     n_epochs_burnin=0,
     n_hypersteps=10,
     marglik_frequency=1,
@@ -36,8 +37,8 @@ def marglik_training(
     fix_sigma_noise=False,
     progress_bar=False,
     enable_backprop=False,
-    dict_key_x='input_ids',
-    dict_key_y='labels',
+    dict_key_x="input_ids",
+    dict_key_y="labels",
 ):
     """Marginal-likelihood based training (Algorithm 1 in [1]).
     Optimize model parameters and hyperparameters jointly.
@@ -137,9 +138,9 @@ def marglik_training(
     losses : list
         list of losses (log joints) obtained during training (to monitor convergence)
     """
-    if 'weight_decay' in optimizer_kwargs:
-        warnings.warn('Weight decay is handled and optimized. Will be set to 0.')
-        optimizer_kwargs['weight_decay'] = 0.0
+    if "weight_decay" in optimizer_kwargs:
+        warnings.warn("Weight decay is handled and optimized. Will be set to 0.")
+        optimizer_kwargs["weight_decay"] = 0.0
 
     # get device, data set size N, number of layers H, number of parameters P
     device = parameters_to_vector(model.parameters()).device
@@ -159,11 +160,11 @@ def marglik_training(
     hyperparameters.append(log_prior_prec)
 
     # set up loss (and observation noise hyperparam)
-    if likelihood == 'classification':
-        criterion = CrossEntropyLoss(reduction='mean')
+    if likelihood == "classification":
+        criterion = CrossEntropyLoss(reduction="mean")
         sigma_noise = 1.0
-    elif likelihood == 'regression':
-        criterion = MSELoss(reduction='mean')
+    elif likelihood == "regression":
+        criterion = MSELoss(reduction="mean")
         log_sigma_noise_init = np.log(sigma_noise_init)
         log_sigma_noise = log_sigma_noise_init * torch.ones(1, device=device)
         log_sigma_noise.requires_grad = True
@@ -195,8 +196,8 @@ def marglik_training(
         disable=not progress_bar,
         position=1,
         leave=False,
-        desc='[Training]',
-        colour='blue',
+        desc="[Training]",
+        colour="blue",
     )
     for epoch in pbar:
         epoch_loss = 0
@@ -211,7 +212,7 @@ def marglik_training(
                 X, y = data
                 X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)
             optimizer.zero_grad()
-            if likelihood == 'regression':
+            if likelihood == "regression":
                 sigma_noise = (
                     torch.exp(log_sigma_noise).detach()
                     if not fix_sigma_noise
@@ -230,7 +231,7 @@ def marglik_training(
             loss.backward()
             optimizer.step()
             epoch_loss += loss.cpu().item() * len(y)
-            if likelihood == 'regression':
+            if likelihood == "regression":
                 epoch_perf += (f.detach() - y).square().sum()
             else:
                 epoch_perf += torch.sum(torch.argmax(f.detach(), dim=-1) == y).item()
@@ -241,8 +242,8 @@ def marglik_training(
 
         # compute validation error to report during training
         logging.info(
-            f'MARGLIK[epoch={epoch}]: network training. Loss={losses[-1]:.3f}.'
-            + f'Perf={epoch_perf/N:.3f}'
+            f"MARGLIK[epoch={epoch}]: network training. Loss={losses[-1]:.3f}."
+            + f"Perf={epoch_perf/N:.3f}"
         )
 
         # only update hyperparameters every marglik_frequency steps after burnin
@@ -251,7 +252,7 @@ def marglik_training(
 
         # optimizer hyperparameters by differentiating marglik
         # 1. fit laplace approximation
-        if likelihood == 'classification':
+        if likelihood == "classification":
             sigma_noise = 1
         else:
             sigma_noise = (
@@ -266,7 +267,7 @@ def marglik_training(
             prior_precision=prior_prec,
             temperature=temperature,
             backend=backend,
-            subset_of_weights='all',
+            subset_of_weights="all",
             dict_key_x=dict_key_x,
             dict_key_y=dict_key_y,
         )
@@ -275,7 +276,7 @@ def marglik_training(
         # 2. differentiate wrt. hyperparameters for n_hypersteps
         for _ in range(n_hypersteps):
             hyper_optimizer.zero_grad()
-            if likelihood == 'classification' or fix_sigma_noise:
+            if likelihood == "classification" or fix_sigma_noise:
                 sigma_noise = None
             else:
                 sigma_noise = torch.exp(log_sigma_noise)
@@ -289,7 +290,7 @@ def marglik_training(
         if margliks[-1] < best_marglik:
             best_model_dict = deepcopy(model.state_dict())
             best_precision = deepcopy(prior_prec.detach())
-            if likelihood == 'classification':
+            if likelihood == "classification":
                 best_sigma = 1
             else:
                 best_sigma = (
@@ -299,16 +300,16 @@ def marglik_training(
                 )
             best_marglik = margliks[-1]
             logging.info(
-                f'MARGLIK[epoch={epoch}]: marglik optimization. MargLik={best_marglik:.2f}. '
-                + 'Saving new best model.'
+                f"MARGLIK[epoch={epoch}]: marglik optimization. MargLik={best_marglik:.2f}. "
+                + "Saving new best model."
             )
         else:
             logging.info(
-                f'MARGLIK[epoch={epoch}]: marglik optimization. MargLik={margliks[-1]:.2f}.'
-                + f'No improvement over {best_marglik:.2f}'
+                f"MARGLIK[epoch={epoch}]: marglik optimization. MargLik={margliks[-1]:.2f}."
+                + f"No improvement over {best_marglik:.2f}"
             )
 
-    logging.info('MARGLIK: finished training. Recover best model and fit Laplace.')
+    logging.info("MARGLIK: finished training. Recover best model and fit Laplace.")
     if best_model_dict is not None:
         model.load_state_dict(best_model_dict)
         sigma_noise = best_sigma
@@ -321,7 +322,7 @@ def marglik_training(
         prior_precision=prior_prec,
         temperature=temperature,
         backend=backend,
-        subset_of_weights='all',
+        subset_of_weights="all",
         enable_backprop=enable_backprop,
         dict_key_x=dict_key_x,
         dict_key_y=dict_key_y,
diff --git a/laplace/subnetlaplace.py b/laplace/subnetlaplace.py
index d4e61bd..a074095 100644
--- a/laplace/subnetlaplace.py
+++ b/laplace/subnetlaplace.py
@@ -1,11 +1,10 @@
 import torch
 from torch.distributions import MultivariateNormal
 
-from laplace.baselaplace import ParametricLaplace, FullLaplace, DiagLaplace
-from laplace.curvature import GGNInterface, EFInterface
+from laplace.baselaplace import DiagLaplace, FullLaplace, ParametricLaplace
+from laplace.curvature import EFInterface, GGNInterface
 
-
-__all__ = ['SubnetLaplace', 'FullSubnetLaplace', 'DiagSubnetLaplace']
+__all__ = ["SubnetLaplace", "FullSubnetLaplace", "DiagSubnetLaplace"]
 
 
 class SubnetLaplace(ParametricLaplace):
@@ -66,17 +65,38 @@ class SubnetLaplace(ParametricLaplace):
         arguments passed to the backend on initialization, for example to
         set the number of MC samples for stochastic approximations.
     """
-    def __init__(self, model, likelihood, subnetwork_indices, sigma_noise=1., prior_precision=1.,
-                 prior_mean=0., temperature=1., backend=None, backend_kwargs=None, asdl_fisher_kwargs=None):
+
+    def __init__(
+        self,
+        model,
+        likelihood,
+        subnetwork_indices,
+        sigma_noise=1.0,
+        prior_precision=1.0,
+        prior_mean=0.0,
+        temperature=1.0,
+        backend=None,
+        backend_kwargs=None,
+        asdl_fisher_kwargs=None,
+    ):
         if asdl_fisher_kwargs is not None:
-            raise ValueError('Subnetwork Laplace does not support asdl_fisher_kwargs.')
+            raise ValueError("Subnetwork Laplace does not support asdl_fisher_kwargs.")
         self.H = None
-        super().__init__(model, likelihood, sigma_noise=sigma_noise,
-                         prior_precision=prior_precision, prior_mean=prior_mean,
-                         temperature=temperature, backend=backend, backend_kwargs=backend_kwargs)
+        super().__init__(
+            model,
+            likelihood,
+            sigma_noise=sigma_noise,
+            prior_precision=prior_precision,
+            prior_mean=prior_mean,
+            temperature=temperature,
+            backend=backend,
+            backend_kwargs=backend_kwargs,
+        )
         if backend is not None:
-            if not isinstance(backend, GGNInterface) and not isinstance(backend, EFInterface):
-                raise ValueError('SubnetLaplace can only be used with GGN and EF.')
+            if not isinstance(backend, GGNInterface) and not isinstance(
+                backend, EFInterface
+            ):
+                raise ValueError("SubnetLaplace can only be used with GGN and EF.")
         # check validity of subnetwork indices and pass them to backend
         self._check_subnetwork_indices(subnetwork_indices)
         self.backend.subnetwork_indices = subnetwork_indices
@@ -85,19 +105,30 @@ def __init__(self, model, likelihood, subnetwork_indices, sigma_noise=1., prior_
 
     def _check_subnetwork_indices(self, subnetwork_indices):
         """Check that subnetwork indices are valid indices of the vectorized model parameters
-           (i.e. `torch.nn.utils.parameters_to_vector(model.parameters())`).
+        (i.e. `torch.nn.utils.parameters_to_vector(model.parameters())`).
         """
         if subnetwork_indices is None:
-            raise ValueError('Subnetwork indices cannot be None.')
-        elif not ((isinstance(subnetwork_indices, torch.LongTensor) or
-                   isinstance(subnetwork_indices, torch.cuda.LongTensor)) and
-                   subnetwork_indices.numel() > 0 and len(subnetwork_indices.shape) == 1):
-            raise ValueError('Subnetwork indices must be non-empty 1-dimensional torch.LongTensor.')
-        elif not (len(subnetwork_indices[subnetwork_indices < 0]) == 0 and
-            len(subnetwork_indices[subnetwork_indices >= self.n_params]) == 0):
-            raise ValueError(f'Subnetwork indices must lie between 0 and n_params={self.n_params}.')
+            raise ValueError("Subnetwork indices cannot be None.")
+        elif not (
+            (
+                isinstance(subnetwork_indices, torch.LongTensor)
+                or isinstance(subnetwork_indices, torch.cuda.LongTensor)
+            )
+            and subnetwork_indices.numel() > 0
+            and len(subnetwork_indices.shape) == 1
+        ):
+            raise ValueError(
+                "Subnetwork indices must be non-empty 1-dimensional torch.LongTensor."
+            )
+        elif not (
+            len(subnetwork_indices[subnetwork_indices < 0]) == 0
+            and len(subnetwork_indices[subnetwork_indices >= self.n_params]) == 0
+        ):
+            raise ValueError(
+                f"Subnetwork indices must lie between 0 and n_params={self.n_params}."
+            )
         elif not (len(subnetwork_indices.unique()) == len(subnetwork_indices)):
-            raise ValueError('Subnetwork indices must not contain duplicate entries.')
+            raise ValueError("Subnetwork indices must not contain duplicate entries.")
 
     @property
     def prior_precision_diag(self):
@@ -109,13 +140,15 @@ def prior_precision_diag(self):
         prior_precision_diag : torch.Tensor
         """
         if len(self.prior_precision) == 1:  # scalar
-            return self.prior_precision * torch.ones(self.n_params_subnet, device=self._device)
+            return self.prior_precision * torch.ones(
+                self.n_params_subnet, device=self._device
+            )
 
         elif len(self.prior_precision) == self.n_params_subnet:  # diagonal
             return self.prior_precision
 
         else:
-            raise ValueError('Mismatch of prior and model. Diagonal or scalar prior.')
+            raise ValueError("Mismatch of prior and model. Diagonal or scalar prior.")
 
     @property
     def mean_subnet(self):
@@ -123,7 +156,7 @@ def mean_subnet(self):
 
     @property
     def scatter(self):
-        delta = (self.mean_subnet - self.prior_mean)
+        delta = self.mean_subnet - self.prior_mean
         return (delta * self.prior_precision_diag) @ delta
 
     def assemble_full_samples(self, subnet_samples):
@@ -139,11 +172,14 @@ class FullSubnetLaplace(SubnetLaplace, FullLaplace):
     Mathematically, we have \\(P \\in \\mathbb{R}^{P \\times P}\\).
     See `FullLaplace`, `SubnetLaplace`, and `BaseLaplace` for the full interface.
     """
+
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
-    _key = ('subnetwork', 'full')
+    _key = ("subnetwork", "full")
 
     def _init_H(self):
-        self.H = torch.zeros(self.n_params_subnet, self.n_params_subnet, device=self._device)
+        self.H = torch.zeros(
+            self.n_params_subnet, self.n_params_subnet, device=self._device
+        )
 
     def sample(self, n_samples=100, generator=None):
         # sample only subnetwork parameters and set all other parameters to their MAP estimates
@@ -158,24 +194,27 @@ class DiagSubnetLaplace(SubnetLaplace, DiagLaplace):
     Mathematically, we have \\(P \\approx \\textrm{diag}(P)\\).
     See `DiagLaplace`, `SubnetLaplace`, and `BaseLaplace` for the full interface.
     """
+
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
-    _key = ('subnetwork', 'diag')
+    _key = ("subnetwork", "diag")
 
     def _init_H(self):
         self.H = torch.zeros(self.n_params_subnet, device=self._device)
 
     def _check_jacobians(self, Js):
         if not isinstance(Js, torch.Tensor):
-            raise ValueError('Jacobians have to be torch.Tensor.')
+            raise ValueError("Jacobians have to be torch.Tensor.")
         if not Js.device == self._device:
-            raise ValueError('Jacobians need to be on the same device as Laplace.')
+            raise ValueError("Jacobians need to be on the same device as Laplace.")
         m, k, p = Js.size()
         if p != self.n_params_subnet:
-            raise ValueError('Invalid Jacobians shape for Laplace posterior approx.')
+            raise ValueError("Invalid Jacobians shape for Laplace posterior approx.")
 
     def sample(self, n_samples=100, generator=None):
         # sample only subnetwork parameters and set all other parameters to their MAP estimates
-        samples = torch.randn(n_samples, self.n_params_subnet, device=self._device, generator=generator)
+        samples = torch.randn(
+            n_samples, self.n_params_subnet, device=self._device, generator=generator
+        )
         samples = samples * self.posterior_scale.reshape(1, self.n_params_subnet)
         subnet_samples = self.mean_subnet.reshape(1, self.n_params_subnet) + samples
         return self.assemble_full_samples(subnet_samples)
diff --git a/laplace/utils/__init__.py b/laplace/utils/__init__.py
index 9e97911..948250d 100644
--- a/laplace/utils/__init__.py
+++ b/laplace/utils/__init__.py
@@ -1,17 +1,58 @@
-from laplace.utils.utils import get_nll, validate, parameters_per_layer, invsqrt_precision, _is_batchnorm, _is_valid_scalar, kron, diagonal_add_scalar, symeig, block_diag, expand_prior_precision, normal_samples, fix_prior_prec_structure
 from laplace.utils.feature_extractor import FeatureExtractor
 from laplace.utils.matrix import Kron, KronDecomposed
-from laplace.utils.swag import fit_diagonal_swag_var
-from laplace.utils.subnetmask import SubnetMask, RandomSubnetMask, LargestMagnitudeSubnetMask, LargestVarianceDiagLaplaceSubnetMask, LargestVarianceSWAGSubnetMask, ParamNameSubnetMask, ModuleNameSubnetMask, LastLayerSubnetMask
 from laplace.utils.metrics import RunningNLLMetric
+from laplace.utils.subnetmask import (
+    LargestMagnitudeSubnetMask,
+    LargestVarianceDiagLaplaceSubnetMask,
+    LargestVarianceSWAGSubnetMask,
+    LastLayerSubnetMask,
+    ModuleNameSubnetMask,
+    ParamNameSubnetMask,
+    RandomSubnetMask,
+    SubnetMask,
+)
+from laplace.utils.swag import fit_diagonal_swag_var
+from laplace.utils.utils import (
+    _is_batchnorm,
+    _is_valid_scalar,
+    block_diag,
+    diagonal_add_scalar,
+    expand_prior_precision,
+    fix_prior_prec_structure,
+    get_nll,
+    invsqrt_precision,
+    kron,
+    normal_samples,
+    parameters_per_layer,
+    symeig,
+    validate,
+)
 
-
-__all__ = ['get_nll', 'validate', 'parameters_per_layer', 'invsqrt_precision', 'kron',
-		   'diagonal_add_scalar', 'symeig', 'block_diag', 'normal_samples',
-           '_is_batchnorm', '_is_valid_scalar',
-           'expand_prior_precision', 'fix_prior_prec_structure',
-		   'FeatureExtractor',
-           'Kron', 'KronDecomposed',
-		   'fit_diagonal_swag_var',
-		   'SubnetMask', 'RandomSubnetMask', 'LargestMagnitudeSubnetMask', 'LargestVarianceDiagLaplaceSubnetMask',
-		   'LargestVarianceSWAGSubnetMask', 'ParamNameSubnetMask', 'ModuleNameSubnetMask', 'LastLayerSubnetMask', 'RunningNLLMetric']
+__all__ = [
+    "get_nll",
+    "validate",
+    "parameters_per_layer",
+    "invsqrt_precision",
+    "kron",
+    "diagonal_add_scalar",
+    "symeig",
+    "block_diag",
+    "normal_samples",
+    "_is_batchnorm",
+    "_is_valid_scalar",
+    "expand_prior_precision",
+    "fix_prior_prec_structure",
+    "FeatureExtractor",
+    "Kron",
+    "KronDecomposed",
+    "fit_diagonal_swag_var",
+    "SubnetMask",
+    "RandomSubnetMask",
+    "LargestMagnitudeSubnetMask",
+    "LargestVarianceDiagLaplaceSubnetMask",
+    "LargestVarianceSWAGSubnetMask",
+    "ParamNameSubnetMask",
+    "ModuleNameSubnetMask",
+    "LastLayerSubnetMask",
+    "RunningNLLMetric",
+]
diff --git a/laplace/utils/feature_extractor.py b/laplace/utils/feature_extractor.py
index f5f6878..92ec4f8 100644
--- a/laplace/utils/feature_extractor.py
+++ b/laplace/utils/feature_extractor.py
@@ -4,13 +4,13 @@
 import torch
 import torch.nn as nn
 
-__all__ = ['FeatureExtractor']
+__all__ = ["FeatureExtractor"]
 
 
 class FeatureReduction(str, Enum):
-    PICK_FIRST = 'pick_first'
-    PICK_LAST = 'pick_last'
-    AVERAGE = 'average'
+    PICK_FIRST = "pick_first"
+    PICK_LAST = "pick_last"
+    AVERAGE = "average"
 
 
 class FeatureExtractor(nn.Module):
@@ -54,7 +54,7 @@ def __init__(
             fr.value for fr in FeatureReduction
         ]:
             raise ValueError(
-                '`feature_reduction` must take value in the `FeatureReduction enum` or '
+                "`feature_reduction` must take value in the `FeatureReduction enum` or "
                 "one of `{'pick_first', 'pick_last', 'average'}`!"
             )
 
@@ -105,9 +105,13 @@ def forward_with_features(
             n_intermediate_dims = len(features.shape) - 2
 
             if self.feature_reduction == FeatureReduction.PICK_FIRST:
-                features = features[:, *([0] * n_intermediate_dims), :].squeeze()
+                features = features[
+                    (slice(None), *([0] * n_intermediate_dims), slice(None))
+                ].squeeze()
             elif self.feature_reduction == FeatureReduction.PICK_LAST:
-                features = features[:, *([0] * n_intermediate_dims), :].squeeze()
+                features = features[
+                    (slice(None), *([-1] * n_intermediate_dims), slice(None))
+                ].squeeze()
             else:
                 ndim = features.ndim
                 features = features.mean(
@@ -129,7 +133,7 @@ def set_last_layer(self, last_layer_name: str) -> None:
         self._last_layer_name = last_layer_name
         self.last_layer = dict(self.model.named_modules())[last_layer_name]
         if not isinstance(self.last_layer, nn.Linear):
-            raise ValueError('Use model with a linear last layer.')
+            raise ValueError("Use model with a linear last layer.")
 
         # set forward hook to extract features in future forward passes
         self.last_layer.register_forward_hook(self._get_hook(last_layer_name))
@@ -158,7 +162,7 @@ def find_last_layer(self, x: torch.Tensor) -> torch.Tensor:
             one batch of data to use as input for the forward pass
         """
         if self.last_layer is not None:
-            raise ValueError('Last layer is already known.')
+            raise ValueError("Last layer is already known.")
 
         act_out = dict()
 
@@ -182,7 +186,7 @@ def act_hook(_, input, __):
         # check if model has more than one module
         # (there might be pathological exceptions)
         if len(handles) <= 2:
-            raise ValueError('The model only has one module.')
+            raise ValueError("The model only has one module.")
 
         # forward pass to find execution order
         out = self.model(x)
@@ -199,4 +203,4 @@ def act_hook(_, input, __):
 
                 return out
 
-        raise ValueError('Something went wrong (all modules have children).')
+        raise ValueError("Something went wrong (all modules have children).")
diff --git a/laplace/utils/matrix.py b/laplace/utils/matrix.py
index e58d805..4b1e301 100644
--- a/laplace/utils/matrix.py
+++ b/laplace/utils/matrix.py
@@ -1,13 +1,13 @@
 from math import pow
-import torch
-import numpy as np
 from typing import Union
-import opt_einsum as oe
 
-from laplace.utils import _is_valid_scalar, symeig, kron, block_diag
+import numpy as np
+import opt_einsum as oe
+import torch
 
+from laplace.utils.utils import _is_valid_scalar, block_diag, kron, symeig
 
-__all__ = ['Kron', 'KronDecomposed']
+__all__ = ["Kron", "KronDecomposed"]
 
 
 class Kron:
@@ -63,7 +63,7 @@ def init_from_model(cls, model, device):
                     ]
                 )
             else:
-                raise ValueError('Invalid parameter shape in network.')
+                raise ValueError("Invalid parameter shape in network.")
         return cls(kfacs)
 
     def __add__(self, other):
@@ -78,7 +78,7 @@ def __add__(self, other):
         kron : Kron
         """
         if not isinstance(other, Kron):
-            raise ValueError('Can only add Kron to Kron.')
+            raise ValueError("Can only add Kron to Kron.")
 
         kfacs = [
             [Hi.add(Hj) for Hi, Hj in zip(Fi, Fj)]
@@ -101,7 +101,7 @@ def __mul__(self, scalar: Union[float, torch.Tensor]):
         kron : Kron
         """
         if not _is_valid_scalar(scalar):
-            raise ValueError('Input not valid python or torch scalar.')
+            raise ValueError("Input not valid python or torch scalar.")
 
         # distribute factors evenly so that each group is multiplied by factor
         kfacs = [[pow(scalar, 1 / len(F)) * Hi for Hi in F] for F in self.kfacs]
@@ -127,14 +127,14 @@ def decompose(self, damping=False):
             for Hi in F:
                 if Hi.ndim > 1:
                     # Dense Kronecker factor.
-                    l, Q = symeig(Hi)
+                    eigval, Q = symeig(Hi)
                 else:
                     # Diagonal Kronecker factor.
-                    l = Hi
+                    eigval = Hi
                     # This might be too memory intensive since len(Hi) can be large.
                     Q = torch.eye(len(Hi), dtype=Hi.dtype, device=Hi.device)
                 Qs.append(Q)
-                ls.append(l)
+                ls.append(eigval)
             eigvecs.append(Qs)
             eigvals.append(ls)
         return KronDecomposed(eigvecs, eigvals, damping=damping)
@@ -175,7 +175,7 @@ def _bmm(self, W: torch.Tensor) -> torch.Tensor:
                 SW.append(QW_pHt.reshape(B * K, p_in * p_out))
                 cur_p += p
             else:
-                raise AttributeError('Shape mismatch')
+                raise AttributeError("Shape mismatch")
         SW = torch.cat(SW, dim=1).reshape(B, K, P)
         return SW
 
@@ -199,7 +199,7 @@ def bmm(self, W: torch.Tensor, exponent: float = 1) -> torch.Tensor:
             result `(batch, classes, params)`
         """
         if exponent != 1:
-            raise ValueError('Only supported after decomposition.')
+            raise ValueError("Only supported after decomposition.")
         if W.ndim == 1:
             return self._bmm(W.unsqueeze(0).unsqueeze(0)).squeeze()
         elif W.ndim == 2:
@@ -207,7 +207,7 @@ def bmm(self, W: torch.Tensor, exponent: float = 1) -> torch.Tensor:
         elif W.ndim == 3:
             return self._bmm(W)
         else:
-            raise ValueError('Invalid shape for W')
+            raise ValueError("Invalid shape for W")
 
     def logdet(self) -> torch.Tensor:
         """Compute log determinant of the Kronecker factors and sums them up.
@@ -308,7 +308,7 @@ def detach(self):
 
     def _check_deltas(self, deltas: torch.Tensor):
         if not isinstance(deltas, torch.Tensor):
-            raise ValueError('Can only add torch.Tensor to KronDecomposed.')
+            raise ValueError("Can only add torch.Tensor to KronDecomposed.")
 
         if deltas.ndim == 0 or (  # scalar
             deltas.ndim == 1  # vector of length 1 or len(self)
@@ -316,7 +316,7 @@ def _check_deltas(self, deltas: torch.Tensor):
         ):
             return
         else:
-            raise ValueError('Invalid shape of delta added to KronDecomposed.')
+            raise ValueError("Invalid shape of delta added to KronDecomposed.")
 
     def __add__(self, deltas: torch.Tensor):
         """Add scalar per layer or only scalar to Kronecker factors.
@@ -346,10 +346,11 @@ def __mul__(self, scalar):
         kron : KronDecomposed
         """
         if not _is_valid_scalar(scalar):
-            raise ValueError('Invalid argument, can only multiply Kron with scalar.')
+            raise ValueError("Invalid argument, can only multiply Kron with scalar.")
 
         eigenvalues = [
-            [pow(scalar, 1 / len(ls)) * l for l in ls] for ls in self.eigenvalues
+            [pow(scalar, 1 / len(ls)) * eigval for eigval in ls]
+            for ls in self.eigenvalues
         ]
         return KronDecomposed(self.eigenvectors, eigenvalues, self.deltas)
 
@@ -378,7 +379,7 @@ def logdet(self) -> torch.Tensor:
                 else:
                     logdet += torch.log(torch.outer(l1, l2) + delta).sum()
             else:
-                raise ValueError('Too many Kronecker factors. Something went wrong.')
+                raise ValueError("Too many Kronecker factors. Something went wrong.")
         return logdet
 
     def _bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
@@ -406,8 +407,8 @@ def _bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
             zip(self.eigenvalues, self.eigenvectors, self.deltas)
         ):
             if len(ls) == 1:
-                Q, l, p = Qs[0], ls[0], len(ls[0])
-                ldelta_exp = torch.pow(l + delta, exponent).reshape(-1, 1)
+                Q, eigval, p = Qs[0], ls[0], len(ls[0])
+                ldelta_exp = torch.pow(eigval + delta, exponent).reshape(-1, 1)
                 W_p = W[:, cur_p : cur_p + p].T
                 SW.append((Q @ (ldelta_exp * (Q.T @ W_p))).T)
                 cur_p += p
@@ -429,7 +430,7 @@ def _bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
                 SW.append(W_p.reshape(B * K, p_in * p_out))
                 cur_p += p
             else:
-                raise AttributeError('Shape mismatch')
+                raise AttributeError("Shape mismatch")
         SW = torch.cat(SW, dim=1).reshape(B, K, P)
         return SW
 
@@ -463,7 +464,7 @@ def bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
         elif W.ndim == 3:
             return self._bmm(W, exponent)
         else:
-            raise ValueError('Invalid shape for W')
+            raise ValueError("Invalid shape for W")
 
     def diag(self, exponent: float = 1) -> torch.Tensor:
         """Extract diagonal of the entire decomposed Kronecker factorization.
@@ -482,7 +483,7 @@ def diag(self, exponent: float = 1) -> torch.Tensor:
             if len(ls) == 1:
                 Ql = Qs[0] * torch.pow(ls[0] + delta, exponent).reshape(1, -1)
                 d = torch.einsum(
-                    'mp,mp->m', Ql, Qs[0]
+                    "mp,mp->m", Ql, Qs[0]
                 )  # only compute inner products for diag
                 diags.append(d)
             else:
@@ -490,12 +491,12 @@ def diag(self, exponent: float = 1) -> torch.Tensor:
                 l1, l2 = ls
                 if self.damping:
                     delta_sqrt = torch.sqrt(delta)
-                    l = torch.pow(
+                    eigval = torch.pow(
                         torch.outer(l1 + delta_sqrt, l2 + delta_sqrt), exponent
                     )
                 else:
-                    l = torch.pow(torch.outer(l1, l2) + delta, exponent)
-                d = oe.contract('mp,nq,pq,mp,nq->mn', Q1, Q2, l, Q1, Q2).flatten()
+                    eigval = torch.pow(torch.outer(l1, l2) + delta, exponent)
+                d = oe.contract("mp,nq,pq,mp,nq->mn", Q1, Q2, eigval, Q1, Q2).flatten()
                 diags.append(d)
         return torch.cat(diags)
 
@@ -516,20 +517,20 @@ def to_matrix(self, exponent: float = 1) -> torch.Tensor:
         blocks = list()
         for Qs, ls, delta in zip(self.eigenvectors, self.eigenvalues, self.deltas):
             if len(ls) == 1:
-                Q, l = Qs[0], ls[0]
-                blocks.append(Q @ torch.diag(torch.pow(l + delta, exponent)) @ Q.T)
+                Q, eigval = Qs[0], ls[0]
+                blocks.append(Q @ torch.diag(torch.pow(eigval + delta, exponent)) @ Q.T)
             else:
                 Q1, Q2 = Qs
                 l1, l2 = ls
                 Q = kron(Q1, Q2)
                 if self.damping:
                     delta_sqrt = torch.sqrt(delta)
-                    l = torch.pow(
+                    eigval = torch.pow(
                         torch.outer(l1 + delta_sqrt, l2 + delta_sqrt), exponent
                     )
                 else:
-                    l = torch.pow(torch.outer(l1, l2) + delta, exponent)
-                L = torch.diag(l.flatten())
+                    eigval = torch.pow(torch.outer(l1, l2) + delta, exponent)
+                L = torch.diag(eigval.flatten())
                 blocks.append(Q @ L @ Q.T)
         return block_diag(blocks)
 
diff --git a/laplace/utils/metrics.py b/laplace/utils/metrics.py
index 6a29d3e..254aa7e 100644
--- a/laplace/utils/metrics.py
+++ b/laplace/utils/metrics.py
@@ -12,10 +12,13 @@ class RunningNLLMetric(Metric):
     ignore_index: int, default = -100
         which class label to ignore when computing the NLL loss
     """
+
     def __init__(self, ignore_index=-100):
         super().__init__()
-        self.add_state('nll_sum', default=torch.tensor(0.), dist_reduce_fx='sum')
-        self.add_state('n_valid_labels', default=torch.tensor(0.), dist_reduce_fx='sum')
+        self.add_state("nll_sum", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state(
+            "n_valid_labels", default=torch.tensor(0.0), dist_reduce_fx="sum"
+        )
         self.ignore_index = ignore_index
 
     def update(self, probs: torch.Tensor, targets: torch.Tensor) -> None:
@@ -32,7 +35,7 @@ def update(self, probs: torch.Tensor, targets: torch.Tensor) -> None:
         targets = targets.view(-1)
 
         self.nll_sum += F.nll_loss(
-            probs.log(), targets, ignore_index=self.ignore_index, reduction='sum'
+            probs.log(), targets, ignore_index=self.ignore_index, reduction="sum"
         )
         self.n_valid_labels += (targets != self.ignore_index).sum()
 
diff --git a/laplace/utils/subnetmask.py b/laplace/utils/subnetmask.py
index e03f489..0d19679 100644
--- a/laplace/utils/subnetmask.py
+++ b/laplace/utils/subnetmask.py
@@ -4,12 +4,19 @@
 from torch.nn import CrossEntropyLoss, MSELoss
 from torch.nn.utils import parameters_to_vector
 
-from laplace.utils import FeatureExtractor, fit_diagonal_swag_var
+from laplace.utils.feature_extractor import FeatureExtractor
+from laplace.utils.swag import fit_diagonal_swag_var
 
-
-__all__ = ['SubnetMask', 'RandomSubnetMask', 'LargestMagnitudeSubnetMask',
-           'LargestVarianceDiagLaplaceSubnetMask', 'LargestVarianceSWAGSubnetMask',
-           'ParamNameSubnetMask', 'ModuleNameSubnetMask', 'LastLayerSubnetMask']
+__all__ = [
+    "SubnetMask",
+    "RandomSubnetMask",
+    "LargestMagnitudeSubnetMask",
+    "LargestVarianceDiagLaplaceSubnetMask",
+    "LargestVarianceSWAGSubnetMask",
+    "ParamNameSubnetMask",
+    "ModuleNameSubnetMask",
+    "LastLayerSubnetMask",
+]
 
 
 class SubnetMask:
@@ -19,6 +26,7 @@ class SubnetMask:
     ----------
     model : torch.nn.Module
     """
+
     def __init__(self, model):
         self.model = model
         self.parameter_vector = parameters_to_vector(self.model.parameters()).detach()
@@ -28,7 +36,7 @@ def __init__(self, model):
 
     def _check_select(self):
         if self._indices is None:
-            raise AttributeError('Subnetwork mask not selected. Run select() first.')
+            raise AttributeError("Subnetwork mask not selected. Run select() first.")
 
     @property
     def _device(self):
@@ -64,23 +72,39 @@ def convert_subnet_mask_to_indices(self, subnet_mask):
             that define the subnetwork
         """
         if not isinstance(subnet_mask, torch.Tensor):
-            raise ValueError('Subnetwork mask needs to be torch.Tensor!')
-        elif subnet_mask.dtype not in [torch.int64, torch.int32, torch.int16, torch.int8,
-            torch.uint8, torch.bool] or len(subnet_mask.shape) != 1:
+            raise ValueError("Subnetwork mask needs to be torch.Tensor!")
+        elif (
+            subnet_mask.dtype
+            not in [
+                torch.int64,
+                torch.int32,
+                torch.int16,
+                torch.int8,
+                torch.uint8,
+                torch.bool,
+            ]
+            or len(subnet_mask.shape) != 1
+        ):
+            raise ValueError(
+                "Subnetwork mask needs to be 1-dimensional integral or boolean tensor!"
+            )
+        elif (
+            len(subnet_mask) != self._n_params
+            or len(subnet_mask[subnet_mask == 0]) + len(subnet_mask[subnet_mask == 1])
+            != self._n_params
+        ):
             raise ValueError(
-                'Subnetwork mask needs to be 1-dimensional integral or boolean tensor!')
-        elif (len(subnet_mask) != self._n_params or len(subnet_mask[subnet_mask == 0])
-            + len(subnet_mask[subnet_mask == 1]) != self._n_params):
-            raise ValueError('Subnetwork mask needs to be a binary vector of'
-                             'size (n_params) where 1s locate the subnetwork'
-                             'parameters within the vectorized model parameters'
-                             '(i.e. `torch.nn.utils.parameters_to_vector(model.parameters())`)!')
+                "Subnetwork mask needs to be a binary vector of"
+                "size (n_params) where 1s locate the subnetwork"
+                "parameters within the vectorized model parameters"
+                "(i.e. `torch.nn.utils.parameters_to_vector(model.parameters())`)!"
+            )
 
         subnet_mask_indices = subnet_mask.nonzero(as_tuple=True)[0]
         return subnet_mask_indices
 
     def select(self, train_loader=None):
-        """ Select the subnetwork mask.
+        """Select the subnetwork mask.
 
         Parameters
         ----------
@@ -96,14 +120,14 @@ def select(self, train_loader=None):
             that define the subnetwork
         """
         if self._indices is not None:
-            raise ValueError('Subnetwork mask already selected.')
+            raise ValueError("Subnetwork mask already selected.")
 
         subnet_mask = self.get_subnet_mask(train_loader)
         self._indices = self.convert_subnet_mask_to_indices(subnet_mask)
         return self._indices
 
     def get_subnet_mask(self, train_loader):
-        """ Get the subnetwork mask.
+        """Get the subnetwork mask.
 
         Parameters
         ----------
@@ -131,15 +155,18 @@ class ScoreBasedSubnetMask(SubnetMask):
     n_params_subnet : int
         number of parameters in the subnetwork (i.e. number of top-scoring parameters to select)
     """
+
     def __init__(self, model, n_params_subnet):
         super().__init__(model)
 
         if n_params_subnet is None:
             raise ValueError(
-                'Need to pass number of subnetwork parameters when using subnetwork Laplace.')
+                "Need to pass number of subnetwork parameters when using subnetwork Laplace."
+            )
         if n_params_subnet > self._n_params:
             raise ValueError(
-                f'Subnetwork ({n_params_subnet}) cannot be larger than model ({self._n_params}).')
+                f"Subnetwork ({n_params_subnet}) cannot be larger than model ({self._n_params})."
+            )
         self._n_params_subnet = n_params_subnet
         self._param_scores = None
 
@@ -148,16 +175,20 @@ def compute_param_scores(self, train_loader):
 
     def _check_param_scores(self):
         if self._param_scores.shape != self.parameter_vector.shape:
-            raise ValueError('Parameter scores need to be of same shape as parameter vector.')
+            raise ValueError(
+                "Parameter scores need to be of same shape as parameter vector."
+            )
 
     def get_subnet_mask(self, train_loader):
-        """ Get the subnetwork mask by (descendingly) ranking parameters based on their scores."""
+        """Get the subnetwork mask by (descendingly) ranking parameters based on their scores."""
 
         if self._param_scores is None:
             self._param_scores = self.compute_param_scores(train_loader)
         self._check_param_scores()
 
-        idx = torch.argsort(self._param_scores, descending=True)[:self._n_params_subnet]
+        idx = torch.argsort(self._param_scores, descending=True)[
+            : self._n_params_subnet
+        ]
         idx = idx.sort()[0]
         subnet_mask = torch.zeros_like(self.parameter_vector).bool()
         subnet_mask[idx] = 1
@@ -166,12 +197,14 @@ def get_subnet_mask(self, train_loader):
 
 class RandomSubnetMask(ScoreBasedSubnetMask):
     """Subnetwork mask of parameters sampled uniformly at random."""
+
     def compute_param_scores(self, train_loader):
         return torch.rand_like(self.parameter_vector)
 
 
 class LargestMagnitudeSubnetMask(ScoreBasedSubnetMask):
-    """Subnetwork mask identifying the parameters with the largest magnitude. """
+    """Subnetwork mask identifying the parameters with the largest magnitude."""
+
     def compute_param_scores(self, train_loader):
         return self.parameter_vector.abs()
 
@@ -188,13 +221,14 @@ class LargestVarianceDiagLaplaceSubnetMask(ScoreBasedSubnetMask):
     diag_laplace_model : `laplace.baselaplace.DiagLaplace`
         diagonal Laplace model to use for variance estimation
     """
+
     def __init__(self, model, n_params_subnet, diag_laplace_model):
         super().__init__(model, n_params_subnet)
         self.diag_laplace_model = diag_laplace_model
 
     def compute_param_scores(self, train_loader):
         if train_loader is None:
-            raise ValueError('Need to pass train loader for subnet selection.')
+            raise ValueError("Need to pass train loader for subnet selection.")
 
         self.diag_laplace_model.fit(train_loader)
         return self.diag_laplace_model.posterior_variance
@@ -218,8 +252,16 @@ class LargestVarianceSWAGSubnetMask(ScoreBasedSubnetMask):
     swag_lr : float
         learning rate for SWAG snapshot collection
     """
-    def __init__(self, model, n_params_subnet, likelihood='classification',
-                 swag_n_snapshots=40, swag_snapshot_freq=1, swag_lr=0.01):
+
+    def __init__(
+        self,
+        model,
+        n_params_subnet,
+        likelihood="classification",
+        swag_n_snapshots=40,
+        swag_snapshot_freq=1,
+        swag_lr=0.01,
+    ):
         super().__init__(model, n_params_subnet)
         self.likelihood = likelihood
         self.swag_n_snapshots = swag_n_snapshots
@@ -228,16 +270,20 @@ def __init__(self, model, n_params_subnet, likelihood='classification',
 
     def compute_param_scores(self, train_loader):
         if train_loader is None:
-            raise ValueError('Need to pass train loader for subnet selection.')
-
-        if self.likelihood == 'classification':
-            criterion = CrossEntropyLoss(reduction='mean')
-        elif self.likelihood == 'regression':
-            criterion = MSELoss(reduction='mean')
-        param_variances = fit_diagonal_swag_var(self.model, train_loader, criterion,
-                                                n_snapshots_total=self.swag_n_snapshots,
-                                                snapshot_freq=self.swag_snapshot_freq,
-                                                lr=self.swag_lr)
+            raise ValueError("Need to pass train loader for subnet selection.")
+
+        if self.likelihood == "classification":
+            criterion = CrossEntropyLoss(reduction="mean")
+        elif self.likelihood == "regression":
+            criterion = MSELoss(reduction="mean")
+        param_variances = fit_diagonal_swag_var(
+            self.model,
+            train_loader,
+            criterion,
+            n_snapshots_total=self.swag_n_snapshots,
+            snapshot_freq=self.swag_snapshot_freq,
+            lr=self.swag_lr,
+        )
         return param_variances
 
 
@@ -251,6 +297,7 @@ class ParamNameSubnetMask(SubnetMask):
         list of names of the parameters (as in `model.named_parameters()`)
         that define the subnetwork
     """
+
     def __init__(self, model, parameter_names):
         super().__init__(model)
         self._parameter_names = parameter_names
@@ -259,16 +306,16 @@ def __init__(self, model, parameter_names):
     def _check_param_names(self):
         param_names = deepcopy(self._parameter_names)
         if len(param_names) == 0:
-            raise ValueError(f'Parameter name list cannot be empty.')
+            raise ValueError("Parameter name list cannot be empty.")
 
         for name, _ in self.model.named_parameters():
             if name in param_names:
                 param_names.remove(name)
         if len(param_names) > 0:
-            raise ValueError(f'Parameters {param_names} do not exist in model.')
+            raise ValueError(f"Parameters {param_names} do not exist in model.")
 
     def get_subnet_mask(self, train_loader):
-        """ Get the subnetwork mask identifying the specified parameters."""
+        """Get the subnetwork mask identifying the specified parameters."""
 
         self._check_param_names()
 
@@ -293,6 +340,7 @@ class ModuleNameSubnetMask(SubnetMask):
         list of names of the modules (as in `model.named_modules()`) that define the subnetwork;
         the modules cannot have children, i.e. need to be leaf modules
     """
+
     def __init__(self, model, module_names):
         super().__init__(model)
         self._module_names = module_names
@@ -301,21 +349,23 @@ def __init__(self, model, module_names):
     def _check_module_names(self):
         module_names = deepcopy(self._module_names)
         if len(module_names) == 0:
-            raise ValueError(f'Module name list cannot be empty.')
+            raise ValueError("Module name list cannot be empty.")
 
         for name, module in self.model.named_modules():
             if name in module_names:
                 if len(list(module.children())) > 0:
-                    raise ValueError(f'Module "{name}" has children, which is not supported.')
+                    raise ValueError(
+                        f'Module "{name}" has children, which is not supported.'
+                    )
                 elif len(list(module.parameters())) == 0:
                     raise ValueError(f'Module "{name}" does not have any parameters.')
                 else:
                     module_names.remove(name)
         if len(module_names) > 0:
-            raise ValueError(f'Modules {module_names} do not exist in model.')
+            raise ValueError(f"Modules {module_names} do not exist in model.")
 
     def get_subnet_mask(self, train_loader):
-        """ Get the subnetwork mask identifying the specified modules."""
+        """Get the subnetwork mask identifying the specified modules."""
 
         self._check_module_names()
 
@@ -327,7 +377,9 @@ def get_subnet_mask(self, train_loader):
                 mask_method = torch.ones_like
             else:
                 mask_method = torch.zeros_like
-            subnet_mask_list.append(mask_method(parameters_to_vector(module.parameters())))
+            subnet_mask_list.append(
+                mask_method(parameters_to_vector(module.parameters()))
+            )
         subnet_mask = torch.cat(subnet_mask_list).bool()
         return subnet_mask
 
@@ -341,16 +393,19 @@ class LastLayerSubnetMask(ModuleNameSubnetMask):
     last_layer_name: str, default=None
         name of the model's last layer, if None it will be determined automatically
     """
+
     def __init__(self, model, last_layer_name=None):
         super().__init__(model, None)
-        self._feature_extractor = FeatureExtractor(self.model, last_layer_name=last_layer_name)
+        self._feature_extractor = FeatureExtractor(
+            self.model, last_layer_name=last_layer_name
+        )
         self._n_params_subnet = None
 
     def get_subnet_mask(self, train_loader):
-        """ Get the subnetwork mask identifying the last layer."""
+        """Get the subnetwork mask identifying the last layer."""
 
         if train_loader is None:
-            raise ValueError('Need to pass train loader for subnet selection.')
+            raise ValueError("Need to pass train loader for subnet selection.")
 
         self._feature_extractor.eval()
         if self._feature_extractor.last_layer is None:
diff --git a/laplace/utils/swag.py b/laplace/utils/swag.py
index a6aba70..5bac456 100644
--- a/laplace/utils/swag.py
+++ b/laplace/utils/swag.py
@@ -3,28 +3,36 @@
 import torch
 from torch.nn.utils import parameters_to_vector
 
-
-__all__ = ['fit_diagonal_swag_var']
+__all__ = ["fit_diagonal_swag_var"]
 
 
 def _param_vector(model):
     return parameters_to_vector(model.parameters()).detach()
 
 
-def fit_diagonal_swag_var(model, train_loader, criterion, n_snapshots_total=40, snapshot_freq=1,
-                          lr=0.01, momentum=0.9, weight_decay=3e-4, min_var=1e-30):
+def fit_diagonal_swag_var(
+    model,
+    train_loader,
+    criterion,
+    n_snapshots_total=40,
+    snapshot_freq=1,
+    lr=0.01,
+    momentum=0.9,
+    weight_decay=3e-4,
+    min_var=1e-30,
+):
     """
     Fit diagonal SWAG [1], which estimates marginal variances of model parameters by
     computing the first and second moment of SGD iterates with a large learning rate.
-    
+
     Implementation partly adapted from:
     - https://github.com/wjmaddox/swa_gaussian/blob/master/swag/posteriors/swag.py
     - https://github.com/wjmaddox/swa_gaussian/blob/master/experiments/train/run_swag.py
 
     References
     ----------
-    [1] Maddox, W., Garipov, T., Izmailov, P., Vetrov, D., Wilson, AG. 
-    [*A Simple Baseline for Bayesian Uncertainty in Deep Learning*](https://arxiv.org/abs/1902.02476). 
+    [1] Maddox, W., Garipov, T., Izmailov, P., Vetrov, D., Wilson, AG.
+    [*A Simple Baseline for Bayesian Uncertainty in Deep Learning*](https://arxiv.org/abs/1902.02476).
     NeurIPS 2019.
 
     Parameters
@@ -65,7 +73,8 @@ def fit_diagonal_swag_var(model, train_loader, criterion, n_snapshots_total=40,
 
     # run SGD to collect model snapshots
     optimizer = torch.optim.SGD(
-        _model.parameters(), lr=lr, momentum=momentum, weight_decay=weight_decay)
+        _model.parameters(), lr=lr, momentum=momentum, weight_decay=weight_decay
+    )
     n_epochs = snapshot_freq * n_snapshots_total
     for epoch in range(n_epochs):
         for inputs, targets in train_loader:
@@ -83,5 +92,5 @@ def fit_diagonal_swag_var(model, train_loader, criterion, n_snapshots_total=40,
             n_snapshots += 1
 
     # compute marginal parameter variances, Var[P] = E[P^2] - E[P]^2
-    param_variances = torch.clamp(sq_mean - mean ** 2, min_var)
+    param_variances = torch.clamp(sq_mean - mean**2, min_var)
     return param_variances
diff --git a/laplace/utils/utils.py b/laplace/utils/utils.py
index f129109..745c616 100644
--- a/laplace/utils/utils.py
+++ b/laplace/utils/utils.py
@@ -11,15 +11,15 @@
 from torchmetrics import Metric
 
 __all__ = [
-    'get_nll',
-    'validate',
-    'parameters_per_layer',
-    'invsqrt_precision',
-    'kron',
-    'diagonal_add_scalar',
-    'symeig',
-    'block_diag',
-    'expand_prior_precision',
+    "get_nll",
+    "validate",
+    "parameters_per_layer",
+    "invsqrt_precision",
+    "kron",
+    "diagonal_add_scalar",
+    "symeig",
+    "block_diag",
+    "expand_prior_precision",
 ]
 
 
@@ -32,11 +32,11 @@ def validate(
     laplace,
     val_loader,
     loss,
-    pred_type='glm',
-    link_approx='probit',
+    pred_type="glm",
+    link_approx="probit",
     n_samples=100,
     loss_with_var=False,
-    dict_key_y='labels',
+    dict_key_y="labels",
 ) -> float:
     laplace.model.eval()
     assert callable(loss) or isinstance(loss, Metric)
@@ -167,7 +167,7 @@ def diagonal_add_scalar(X, value):
     -------
     X_add_scalar : torch.Tensor
     """
-    if not X.device == torch.device('cpu'):
+    if not X.device == torch.device("cpu"):
         indices = torch.cuda.LongTensor([[i, i] for i in range(X.shape[0])])
     else:
         indices = torch.LongTensor([[i, i] for i in range(X.shape[0])])
@@ -191,18 +191,18 @@ def symeig(M):
         eigenvectors
     """
     try:
-        L, W = torch.linalg.eigh(M, UPLO='U')
+        L, W = torch.linalg.eigh(M, UPLO="U")
     except RuntimeError:  # did not converge
-        logging.info('SYMEIG: adding jitter, did not converge.')
+        logging.info("SYMEIG: adding jitter, did not converge.")
         # use W L W^T + I = W (L + I) W^T
         M = M + torch.eye(M.shape[0], device=M.device)
         try:
-            L, W = torch.linalg.eigh(M, UPLO='U')
+            L, W = torch.linalg.eigh(M, UPLO="U")
             L -= 1.0
         except RuntimeError:
-            stats = f'diag: {M.diagonal()}, max: {M.abs().max()}, '
-            stats = stats + f'min: {M.abs().min()}, mean: {M.abs().mean()}'
-            logging.info(f'SYMEIG: adding jitter failed. Stats: {stats}')
+            stats = f"diag: {M.diagonal()}, max: {M.abs().max()}, "
+            stats = stats + f"min: {M.abs().min()}, mean: {M.abs().mean()}"
+            logging.info(f"SYMEIG: adding jitter failed. Stats: {stats}")
             exit()
     # eigenvalues of symeig at least 0
     L = L.clamp(min=0.0)
@@ -267,14 +267,14 @@ def expand_prior_precision(prior_prec, model):
 def fix_prior_prec_structure(
     prior_prec_init, prior_structure, n_layers, n_params, device
 ):
-    if prior_structure == 'scalar':
+    if prior_structure == "scalar":
         prior_prec_init = torch.full((1,), prior_prec_init, device=device)
-    elif prior_structure == 'layerwise':
+    elif prior_structure == "layerwise":
         prior_prec_init = torch.full((n_layers,), prior_prec_init, device=device)
-    elif prior_structure == 'diag':
+    elif prior_structure == "diag":
         prior_prec_init = torch.full((n_params,), prior_prec_init, device=device)
     else:
-        raise ValueError(f'Invalid prior structure {prior_structure}.')
+        raise ValueError(f"Invalid prior structure {prior_structure}.")
     return prior_prec_init
 
 
@@ -292,7 +292,7 @@ def normal_samples(mean, var, n_samples, generator=None):
     generator : torch.Generator
         random number generator
     """
-    assert mean.ndim == 2, 'Invalid input shape of mean, should be 2-dimensional.'
+    assert mean.ndim == 2, "Invalid input shape of mean, should be 2-dimensional."
     _, output_dim = mean.shape
     randn_samples = torch.randn(
         (output_dim, n_samples),
@@ -313,4 +313,4 @@ def normal_samples(mean, var, n_samples, generator=None):
         )  # expand batch dim
         return (mean.unsqueeze(-1) + scaled_samples).permute((2, 0, 1))
     else:
-        raise ValueError('Invalid input shapes.')
+        raise ValueError("Invalid input shapes.")
diff --git a/pyproject.toml b/pyproject.toml
index 6f1f38a..f8c3afc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,9 +1,16 @@
 [build-system]
-requires = [
-    "setuptools>=42",
-    "wheel"
-]
+requires = ["setuptools>=42", "wheel"]
 build-backend = "setuptools.build_meta"
 
+[tool.ruff.lint]
+# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`)  codes by default.
+# Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or
+# McCabe complexity (`C901`) by default.
+# Adds isort ("I") and pylint errors and warning ("PLE", "PLW") to defaults.
+select = ["E4", "E7", "E9", "F", "I", "PLE", "PLW"]
+# Ignore pylint's overwritten loop variable by assignment error
+# Skip "Module level import not at top of file"
+ignore = ["PLW2901", "E402"]
+
 [tool.ruff.format]
-quote-style = "single"
+quote-style = "double"


Test patch:
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
new file mode 100644
index 0000000..1c49f7f
--- /dev/null
+++ b/.github/workflows/pytest.yml
@@ -0,0 +1,28 @@
+name: pytest
+
+on: pull_request
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        # This is the version of the action for setting up Python, not the Python version.
+        uses: actions/setup-python@v4
+        with:
+          # Semantic version range syntax or exact version of a Python version
+          python-version: "3.9"
+          # Optional - x64 or x86 architecture, defaults to x64
+          architecture: "x64"
+          # cache: 'pip'
+      # You can test your matrix by printing the current Python version
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+          pip install -e .
+      - name: Test with pytest
+        run: |
+          pip install pytest
+          pytest tests
diff --git a/tests/test_baselaplace.py b/tests/test_baselaplace.py
index 1664d59..2d32c3e 100644
--- a/tests/test_baselaplace.py
+++ b/tests/test_baselaplace.py
@@ -1,25 +1,27 @@
+from __future__ import annotations
+
 from collections.abc import MutableMapping
-from math import sqrt, prod
-import pytest
+from copy import deepcopy
 from itertools import product
+from math import prod, sqrt
+
 import numpy as np
-from copy import deepcopy
+import pytest
 import torch
 from torch import nn
+from torch.distributions import Categorical, Normal
 from torch.distributions.multivariate_normal import MultivariateNormal
 from torch.nn.utils import parameters_to_vector
 from torch.utils.data import DataLoader, TensorDataset
-from torch.distributions import Normal, Categorical
-from laplace.curvature.backpack import BackPackEF
-from laplace.curvature.curvlinops import CurvlinopsEF, CurvlinopsGGN
 from torchvision.models import wide_resnet50_2
 
-from laplace.laplace import FullLaplace, KronLaplace, DiagLaplace, LowRankLaplace
+from laplace.curvature import AsdlEF, AsdlGGN, BackPackGGN
+from laplace.curvature.backpack import BackPackEF
+from laplace.curvature.curvlinops import CurvlinopsEF, CurvlinopsGGN
+from laplace.laplace import DiagLaplace, FullLaplace, KronLaplace, LowRankLaplace
 from laplace.utils import KronDecomposed
-from laplace.curvature import AsdlGGN, BackPackGGN, AsdlEF
 from tests.utils import ListDataset, dict_data_collator, jacobians_naive
 
-
 torch.manual_seed(240)
 torch.set_default_tensor_type(torch.DoubleTensor)
 flavors = [FullLaplace, KronLaplace, DiagLaplace, LowRankLaplace]
@@ -29,10 +31,10 @@
 @pytest.fixture
 def model():
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -51,7 +53,7 @@ def __init__(self):
 
         def forward(self, data: MutableMapping | torch.Tensor):
             if isinstance(data, MutableMapping):
-                x = data['test_input_key'].to(next(self.parameters()).device)
+                x = data["test_input_key"].to(next(self.parameters()).device)
             else:
                 x = data
 
@@ -111,8 +113,8 @@ def custom_loader_clf():
     data = []
     for _ in range(10):
         datum = {
-            'test_input_key': torch.randn(5),
-            'test_label_key': torch.randint(2, (1,)),
+            "test_input_key": torch.randn(5),
+            "test_label_key": torch.randint(2, (1,)),
         }
         data.append(datum)
     return DataLoader(ListDataset(data), batch_size=3, collate_fn=dict_data_collator)
@@ -123,8 +125,8 @@ def custom_loader_reg():
     data = []
     for _ in range(10):
         datum = {
-            'test_input_key': torch.randn(5),
-            'test_label_key': torch.randn(2),
+            "test_input_key": torch.randn(5),
+            "test_label_key": torch.randn(2),
         }
         data.append(datum)
     return DataLoader(ListDataset(data), batch_size=3, collate_fn=dict_data_collator)
@@ -143,9 +145,9 @@ def reward_test_X():
     return X
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init(laplace, model):
-    lap = laplace(model, 'classification')
+    lap = laplace(model, "classification")
     assert torch.allclose(lap.mean, lap.prior_mean)
     if laplace in [FullLaplace, DiagLaplace]:
         H = lap.H.clone()
@@ -161,101 +163,101 @@ def test_laplace_init(laplace, model):
                 assert torch.allclose(k1, k2)
 
 
-@pytest.mark.xfail(strict=True)
+@pytest.mark.skip(reason="Does not work well with Github actions")
 def test_laplace_large_init(large_model):
-    lap = FullLaplace(large_model, 'classification')
+    FullLaplace(large_model, "classification")
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_invalid_likelihood(laplace, model):
     with pytest.raises(ValueError):
-        lap = laplace(model, 'otherlh')
+        laplace(model, "otherlh")
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init_noise(laplace, model):
     # float
     sigma_noise = 1.2
-    lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise)
+    laplace(model, likelihood="regression", sigma_noise=sigma_noise)
     # torch.tensor 0-dim
     sigma_noise = torch.tensor(1.2)
-    lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise)
+    laplace(model, likelihood="regression", sigma_noise=sigma_noise)
     # torch.tensor 1-dim
     sigma_noise = torch.tensor(1.2).reshape(-1)
-    lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise)
+    laplace(model, likelihood="regression", sigma_noise=sigma_noise)
 
     # for classification should fail
     sigma_noise = 1.2
     with pytest.raises(ValueError):
-        lap = laplace(model, likelihood='classification', sigma_noise=sigma_noise)
+        laplace(model, likelihood="classification", sigma_noise=sigma_noise)
 
     # other than that should fail
     # higher dim
     sigma_noise = torch.tensor(1.2).reshape(1, 1)
     with pytest.raises(ValueError):
-        lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise)
+        laplace(model, likelihood="regression", sigma_noise=sigma_noise)
     # other datatype, only reals supported
-    sigma_noise = '1.2'
+    sigma_noise = "1.2"
     with pytest.raises(ValueError):
-        lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise)
+        laplace(model, likelihood="regression", sigma_noise=sigma_noise)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init_precision(laplace, model):
     # float
     precision = 10.6
-    lap = laplace(model, likelihood='regression', prior_precision=precision)
+    laplace(model, likelihood="regression", prior_precision=precision)
     # torch.tensor 0-dim
     precision = torch.tensor(10.6)
-    lap = laplace(model, likelihood='regression', prior_precision=precision)
+    laplace(model, likelihood="regression", prior_precision=precision)
     # torch.tensor 1-dim
     precision = torch.tensor(10.7).reshape(-1)
-    lap = laplace(model, likelihood='regression', prior_precision=precision)
+    laplace(model, likelihood="regression", prior_precision=precision)
     # torch.tensor 1-dim param-shape
     precision = torch.tensor(10.7).reshape(-1).repeat(model.n_params)
     if laplace == KronLaplace:
         # Kron should not accept per parameter prior precision
         with pytest.raises(ValueError):
-            lap = laplace(model, likelihood='regression', prior_precision=precision)
+            laplace(model, likelihood="regression", prior_precision=precision)
     else:
-        lap = laplace(model, likelihood='regression', prior_precision=precision)
+        laplace(model, likelihood="regression", prior_precision=precision)
     # torch.tensor 1-dim layer-shape
     precision = torch.tensor(10.7).reshape(-1).repeat(model.n_layers)
-    lap = laplace(model, likelihood='regression', prior_precision=precision)
+    laplace(model, likelihood="regression", prior_precision=precision)
 
     # other than that should fail
     # higher dim
     precision = torch.tensor(10.6).reshape(1, 1)
     with pytest.raises(ValueError):
-        lap = laplace(model, likelihood='regression', prior_precision=precision)
+        laplace(model, likelihood="regression", prior_precision=precision)
     # unmatched dim
     precision = torch.tensor(10.6).reshape(-1).repeat(17)
     with pytest.raises(ValueError):
-        lap = laplace(model, likelihood='regression', prior_precision=precision)
+        laplace(model, likelihood="regression", prior_precision=precision)
     # other datatype, only reals supported
-    precision = '1.5'
+    precision = "1.5"
     with pytest.raises(ValueError):
-        lap = laplace(model, likelihood='regression', prior_precision=precision)
+        laplace(model, likelihood="regression", prior_precision=precision)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init_prior_mean_and_scatter(laplace, model, class_loader):
     mean = parameters_to_vector(model.parameters())
     P = len(mean)
     lap_scalar_mean = laplace(
-        model, 'classification', prior_precision=1e-2, prior_mean=1.0
+        model, "classification", prior_precision=1e-2, prior_mean=1.0
     )
     assert torch.allclose(lap_scalar_mean.prior_mean, torch.tensor([1.0]))
     lap_tensor_mean = laplace(
-        model, 'classification', prior_precision=1e-2, prior_mean=torch.ones(1)
+        model, "classification", prior_precision=1e-2, prior_mean=torch.ones(1)
     )
     assert torch.allclose(lap_tensor_mean.prior_mean, torch.tensor([1.0]))
     lap_tensor_scalar_mean = laplace(
-        model, 'classification', prior_precision=1e-2, prior_mean=torch.ones(1)[0]
+        model, "classification", prior_precision=1e-2, prior_mean=torch.ones(1)[0]
     )
     assert torch.allclose(lap_tensor_scalar_mean.prior_mean, torch.tensor(1.0))
     lap_tensor_full_mean = laplace(
-        model, 'classification', prior_precision=1e-2, prior_mean=torch.ones(P)
+        model, "classification", prior_precision=1e-2, prior_mean=torch.ones(P)
     )
     assert torch.allclose(lap_tensor_full_mean.prior_mean, torch.ones(P))
 
@@ -278,31 +280,31 @@ def test_laplace_init_prior_mean_and_scatter(laplace, model, class_loader):
     # too many dims
     with pytest.raises(ValueError):
         prior_mean = torch.ones(P).unsqueeze(-1)
-        laplace(model, 'classification', prior_precision=1e-2, prior_mean=prior_mean)
+        laplace(model, "classification", prior_precision=1e-2, prior_mean=prior_mean)
 
     # unmatched dim
     with pytest.raises(ValueError):
         prior_mean = torch.ones(P - 3)
-        laplace(model, 'classification', prior_precision=1e-2, prior_mean=prior_mean)
+        laplace(model, "classification", prior_precision=1e-2, prior_mean=prior_mean)
 
     # invalid argument type
     with pytest.raises(ValueError):
-        laplace(model, 'classification', prior_precision=1e-2, prior_mean='72')
+        laplace(model, "classification", prior_precision=1e-2, prior_mean="72")
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init_temperature(laplace, model):
     # valid float
     T = 1.1
-    lap = laplace(model, likelihood='classification', temperature=T)
+    lap = laplace(model, likelihood="classification", temperature=T)
     assert lap.temperature == T
 
 
 @pytest.mark.parametrize(
-    'laplace,lh', product(flavors, ['classification', 'regression'])
+    "laplace,lh", product(flavors, ["classification", "regression"])
 )
 def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
-    if lh == 'classification':
+    if lh == "classification":
         loader = class_loader
         sigma_noise = 1.0
     else:
@@ -319,7 +321,7 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
     # Test log likelihood (Train)
     log_lik = lap.log_likelihood
     # compute true log lik
-    if lh == 'classification':
+    if lh == "classification":
         log_lik_true = Categorical(logits=f).log_prob(y).sum()
         assert torch.allclose(log_lik, log_lik_true)
     else:
@@ -344,8 +346,8 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
     if laplace == DiagLaplace:
         log_det_post_prec = lap.posterior_precision.log().sum()
     elif laplace == LowRankLaplace:
-        (U, l), p0 = lap.posterior_precision
-        log_det_post_prec = (U @ torch.diag(l) @ U.T + p0.diag()).logdet()
+        (U, eigval), p0 = lap.posterior_precision
+        log_det_post_prec = (U @ torch.diag(eigval) @ U.T + p0.diag()).logdet()
     else:
         log_det_post_prec = lap.posterior_precision.logdet()
     lml = lml + 1 / 2 * (prior_prec.logdet() - log_det_post_prec)
@@ -367,19 +369,19 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
     elif laplace == KronLaplace:
         Sigma = lap.posterior_precision.to_matrix(exponent=-1)
     elif laplace == LowRankLaplace:
-        (U, l), p0 = lap.posterior_precision
-        Sigma = (U @ torch.diag(l) @ U.T + p0.diag()).inverse()
+        (U, eigval), p0 = lap.posterior_precision
+        Sigma = (U @ torch.diag(eigval) @ U.T + p0.diag()).inverse()
     elif laplace == DiagLaplace:
         Sigma = torch.diag(lap.posterior_variance)
     Js, f = jacobians_naive(model, loader.dataset.tensors[0])
-    true_f_var = torch.einsum('mkp,pq,mcq->mkc', Js, Sigma, Js)
+    true_f_var = torch.einsum("mkp,pq,mcq->mkc", Js, Sigma, Js)
     comp_f_var = lap.functional_variance(Js)
     assert torch.allclose(true_f_var, comp_f_var, rtol=1e-4)
 
 
-@pytest.mark.parametrize('laplace', online_flavors)
+@pytest.mark.parametrize("laplace", online_flavors)
 def test_overriding_fit(laplace, model, reg_loader):
-    lap = laplace(model, 'regression', sigma_noise=0.3, prior_precision=0.7)
+    lap = laplace(model, "regression", sigma_noise=0.3, prior_precision=0.7)
     lap.fit(reg_loader)
     if type(lap.posterior_precision) is KronDecomposed:
         P = lap.posterior_precision.to_matrix()
@@ -398,9 +400,9 @@ def test_overriding_fit(laplace, model, reg_loader):
     assert lap.n_data == len(reg_loader.dataset)
 
 
-@pytest.mark.parametrize('laplace', online_flavors)
+@pytest.mark.parametrize("laplace", online_flavors)
 def test_online_fit(laplace, model, reg_loader):
-    lap = laplace(model, 'regression', sigma_noise=0.3, prior_precision=0.7)
+    lap = laplace(model, "regression", sigma_noise=0.3, prior_precision=0.7)
     lap.fit(reg_loader)
     if type(lap.H) is KronDecomposed:
         P = lap.H.to_matrix().clone()
@@ -420,7 +422,7 @@ def test_online_fit(laplace, model, reg_loader):
 
 
 def test_log_prob_full(model, class_loader):
-    lap = FullLaplace(model, 'classification', prior_precision=0.7)
+    lap = FullLaplace(model, "classification", prior_precision=0.7)
     theta = torch.randn_like(parameters_to_vector(model.parameters()))
     # posterior without fitting is just prior
     posterior = Normal(loc=torch.zeros_like(theta), scale=sqrt(1 / 0.7))
@@ -433,7 +435,7 @@ def test_log_prob_full(model, class_loader):
 
 
 def test_log_prob_kron(model, class_loader):
-    lap = KronLaplace(model, 'classification', prior_precision=0.24)
+    lap = KronLaplace(model, "classification", prior_precision=0.24)
     theta = torch.randn_like(parameters_to_vector(model.parameters()))
     posterior = Normal(loc=lap.mean, scale=sqrt(1 / 0.24))
     assert torch.allclose(lap.log_prob(theta), posterior.log_prob(theta).sum())
@@ -444,19 +446,19 @@ def test_log_prob_kron(model, class_loader):
     assert torch.allclose(lap.log_prob(theta), posterior.log_prob(theta))
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_regression_predictive(laplace, model, reg_loader):
-    lap = laplace(model, 'regression', sigma_noise=0.3, prior_precision=0.7)
+    lap = laplace(model, "regression", sigma_noise=0.3, prior_precision=0.7)
     lap.fit(reg_loader)
     X, y = reg_loader.dataset.tensors
     f = model(X)
 
     # error
     with pytest.raises(ValueError):
-        lap(X, pred_type='linear')
+        lap(X, pred_type="linear")
 
     # GLM predictive, functional variance tested already above.
-    f_mu_glm, f_var_glm = lap(X, pred_type='glm')
+    f_mu_glm, f_var_glm = lap(X, pred_type="glm")
     assert torch.allclose(f_mu_glm, f)
     assert f_var_glm.shape == torch.Size(
         [f_mu_glm.shape[0], f_mu_glm.shape[1], f_mu_glm.shape[1]]
@@ -464,13 +466,13 @@ def test_regression_predictive(laplace, model, reg_loader):
     assert len(f_mu_glm) == len(X)
 
     # NN predictive (only diagonal variance estimation)
-    f_mu_nn, f_var_nn = lap(X, pred_type='nn', link_approx='mc')
+    f_mu_nn, f_var_nn = lap(X, pred_type="nn", link_approx="mc")
     assert f_mu_nn.shape == f_var_nn.shape
     assert f_var_nn.shape == torch.Size([f_mu_nn.shape[0], f_mu_nn.shape[1]])
     assert len(f_mu_nn) == len(X)
 
     # Test joint prediction
-    f_mu_joint, f_cov_joint = lap(X, pred_type='glm', joint=True)
+    f_mu_joint, f_cov_joint = lap(X, pred_type="glm", joint=True)
     assert len(f_mu_joint.shape) == 1
     assert f_mu_joint.shape[0] == prod(f_mu_glm.shape)
     assert len(f_cov_joint.shape) == 2
@@ -478,133 +480,133 @@ def test_regression_predictive(laplace, model, reg_loader):
 
     # The "diagonal" of the joint cov should equal the non-joint var
     b, k = y.shape
-    f_var_joint = torch.einsum('bkbl->bkl', f_cov_joint.reshape(b, k, b, k))
+    f_var_joint = torch.einsum("bkbl->bkl", f_cov_joint.reshape(b, k, b, k))
     assert torch.allclose(f_var_joint, f_var_glm)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_classification_predictive(laplace, model, class_loader):
-    lap = laplace(model, 'classification', prior_precision=0.7)
+    lap = laplace(model, "classification", prior_precision=0.7)
     lap.fit(class_loader)
     X, y = class_loader.dataset.tensors
     f = torch.softmax(model(X), dim=-1)
 
     # error
     with pytest.raises(ValueError):
-        lap(X, pred_type='linear')
+        lap(X, pred_type="linear")
 
     # GLM predictive
-    f_pred = lap(X, pred_type='glm', link_approx='mc', n_samples=100)
+    f_pred = lap(X, pred_type="glm", link_approx="mc", n_samples=100)
     assert f_pred.shape == f.shape
     assert torch.allclose(
         f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
     )  # sum up to 1
-    f_pred = lap(X, pred_type='glm', link_approx='probit')
+    f_pred = lap(X, pred_type="glm", link_approx="probit")
     assert f_pred.shape == f.shape
     assert torch.allclose(
         f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
     )  # sum up to 1
-    f_pred = lap(X, pred_type='glm', link_approx='bridge')
+    f_pred = lap(X, pred_type="glm", link_approx="bridge")
     assert f_pred.shape == f.shape
     assert torch.allclose(
         f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
     )  # sum up to 1
-    f_pred = lap(X, pred_type='glm', link_approx='bridge_norm')
+    f_pred = lap(X, pred_type="glm", link_approx="bridge_norm")
     assert f_pred.shape == f.shape
     assert torch.allclose(
         f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
     )  # sum up to 1
 
     # NN predictive
-    f_pred = lap(X, pred_type='nn', link_approx='mc', n_samples=100)
+    f_pred = lap(X, pred_type="nn", link_approx="mc", n_samples=100)
     assert f_pred.shape == f.shape
     assert torch.allclose(
         f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
     )  # sum up to 1
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_regression_predictive_samples(laplace, model, reg_loader):
-    lap = laplace(model, 'regression', sigma_noise=0.3, prior_precision=0.7)
+    lap = laplace(model, "regression", sigma_noise=0.3, prior_precision=0.7)
     lap.fit(reg_loader)
     X, y = reg_loader.dataset.tensors
     f = model(X)
 
     # error
     with pytest.raises(ValueError):
-        lap(X, pred_type='linear')
+        lap(X, pred_type="linear")
 
     # GLM predictive, functional variance tested already above.
-    fsamples = lap.predictive_samples(X, pred_type='glm', n_samples=100)
+    fsamples = lap.predictive_samples(X, pred_type="glm", n_samples=100)
     assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
 
     # NN predictive (only diagonal variance estimation)
-    fsamples = lap.predictive_samples(X, pred_type='nn', n_samples=100)
+    fsamples = lap.predictive_samples(X, pred_type="nn", n_samples=100)
     assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_classification_predictive_samples(laplace, model, class_loader):
-    lap = laplace(model, 'classification', prior_precision=0.7)
+    lap = laplace(model, "classification", prior_precision=0.7)
     lap.fit(class_loader)
     X, y = class_loader.dataset.tensors
     f = torch.softmax(model(X), dim=-1)
 
     # error
     with pytest.raises(ValueError):
-        lap(X, pred_type='linear')
+        lap(X, pred_type="linear")
 
     # GLM predictive
-    fsamples = lap.predictive_samples(X, pred_type='glm', n_samples=100)
+    fsamples = lap.predictive_samples(X, pred_type="glm", n_samples=100)
     assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
     assert np.allclose(fsamples.sum().item(), len(f) * 100)  # sum up to 1
 
     # NN predictive
-    fsamples = lap.predictive_samples(X, pred_type='nn', n_samples=100)
+    fsamples = lap.predictive_samples(X, pred_type="nn", n_samples=100)
     assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
     assert np.allclose(fsamples.sum().item(), len(f) * 100)  # sum up to 1
 
 
-@pytest.mark.parametrize('laplace', [KronLaplace, DiagLaplace])
+@pytest.mark.parametrize("laplace", [KronLaplace, DiagLaplace])
 def test_reward_modeling(laplace, reward_model, reward_loader, reward_test_X):
-    lap = laplace(reward_model, 'reward_modeling')
+    lap = laplace(reward_model, "reward_modeling")
     lap.fit(reward_loader)
     f = reward_model(reward_test_X)
 
     # error
     with pytest.raises(ValueError):
-        lap(reward_test_X, pred_type='linear')
+        lap(reward_test_X, pred_type="linear")
 
     # GLM predictive, functional variance tested already above.
-    f_mu, f_var = lap(reward_test_X, pred_type='glm')
+    f_mu, f_var = lap(reward_test_X, pred_type="glm")
     assert torch.allclose(f_mu, f)
     assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1], f_mu.shape[1]])
     assert len(f_mu) == len(reward_test_X)
 
     # NN predictive (only diagonal variance estimation)
-    f_mu, f_var = lap(reward_test_X, pred_type='nn', link_approx='mc')
+    f_mu, f_var = lap(reward_test_X, pred_type="nn", link_approx="mc")
     assert f_mu.shape == f_var.shape
     assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1]])
     assert len(f_mu) == len(reward_test_X)
 
 
-@pytest.mark.parametrize('laplace', [KronLaplace, DiagLaplace])
+@pytest.mark.parametrize("laplace", [KronLaplace, DiagLaplace])
 @pytest.mark.parametrize(
-    'backend', [AsdlEF, AsdlGGN, CurvlinopsEF, CurvlinopsGGN, BackPackGGN, BackPackEF]
+    "backend", [AsdlEF, AsdlGGN, CurvlinopsEF, CurvlinopsGGN, BackPackGGN, BackPackEF]
 )
 @pytest.mark.parametrize(
-    'lik,custom_loader',
+    "lik,custom_loader",
     [
-        ('classification', 'custom_loader_clf'),
-        ('regression', 'custom_loader_reg'),
-        ('reward_modeling', 'custom_loader_clf'),
+        ("classification", "custom_loader_clf"),
+        ("regression", "custom_loader_reg"),
+        ("reward_modeling", "custom_loader_clf"),
     ],
 )
 def test_dict_data(laplace, backend, lik, custom_loader, custom_model, request):
     custom_loader = request.getfixturevalue(custom_loader)
 
     if (
-        'backpack' not in backend.__name__.lower()
+        "backpack" not in backend.__name__.lower()
         and laplace != DiagLaplace
         and laplace != CurvlinopsEF
     ):
@@ -618,11 +620,11 @@ def test_dict_data(laplace, backend, lik, custom_loader, custom_model, request):
         custom_model,
         lik,
         backend=backend,
-        dict_key_x='test_input_key',
-        dict_key_y='test_label_key',
+        dict_key_x="test_input_key",
+        dict_key_y="test_label_key",
     )
 
-    if ('backpack' in backend.__name__.lower()) or (
+    if ("backpack" in backend.__name__.lower()) or (
         laplace == DiagLaplace and backend == CurvlinopsEF
     ):
         # Unsupported, thus raises an exception
@@ -636,40 +638,40 @@ def test_dict_data(laplace, backend, lik, custom_loader, custom_model, request):
     test_data = next(iter(custom_loader))
     f = custom_model(test_data)
 
-    if lik == 'classification':
-        f_pred = lap(test_data, pred_type='glm')
+    if lik == "classification":
+        f_pred = lap(test_data, pred_type="glm")
         assert f_pred.shape == f.shape
         assert torch.allclose(
             f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
         )  # sum up to 1
 
-        f_pred = lap(test_data, pred_type='nn', link_approx='mc')
+        f_pred = lap(test_data, pred_type="nn", link_approx="mc")
         assert f_pred.shape == f.shape
         assert torch.allclose(
             f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
         )
     else:
-        f_pred, f_var = lap(test_data, pred_type='glm')
+        f_pred, f_var = lap(test_data, pred_type="glm")
         assert f_pred.shape == f.shape
         assert torch.allclose(f_pred, f)
         assert f_var.shape == (f_pred.shape[0], f_pred.shape[1], f_pred.shape[1])
 
-        f_pred, f_var = lap(test_data, pred_type='nn', link_approx='mc')
+        f_pred, f_var = lap(test_data, pred_type="nn", link_approx="mc")
         assert f_pred.shape == f.shape
         assert f_var.shape == (f_pred.shape[0], f_pred.shape[1])
 
 
-@pytest.mark.parametrize('laplace', [FullLaplace, KronLaplace, DiagLaplace])
+@pytest.mark.parametrize("laplace", [FullLaplace, KronLaplace, DiagLaplace])
 @pytest.mark.parametrize(
-    'backend', [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
+    "backend", [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
 )
 def test_backprop_glm(laplace, model, reg_loader, backend):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True, backend=backend)
+    lap = laplace(model, "regression", enable_backprop=True, backend=backend)
     lap.fit(reg_loader)
-    f_mu, f_var = lap(X, pred_type='glm')
+    f_mu, f_var = lap(X, pred_type="glm")
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
@@ -681,17 +683,17 @@ def test_backprop_glm(laplace, model, reg_loader, backend):
         assert False
 
 
-@pytest.mark.parametrize('laplace', [FullLaplace, KronLaplace, DiagLaplace])
+@pytest.mark.parametrize("laplace", [FullLaplace, KronLaplace, DiagLaplace])
 @pytest.mark.parametrize(
-    'backend', [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
+    "backend", [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
 )
 def test_backprop_glm_joint(laplace, model, reg_loader, backend):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True, backend=backend)
+    lap = laplace(model, "regression", enable_backprop=True, backend=backend)
     lap.fit(reg_loader)
-    f_mu, f_cov = lap(X, pred_type='glm', joint=True)
+    f_mu, f_cov = lap(X, pred_type="glm", joint=True)
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
@@ -703,17 +705,17 @@ def test_backprop_glm_joint(laplace, model, reg_loader, backend):
         assert False
 
 
-@pytest.mark.parametrize('laplace', [FullLaplace, KronLaplace, DiagLaplace])
+@pytest.mark.parametrize("laplace", [FullLaplace, KronLaplace, DiagLaplace])
 @pytest.mark.parametrize(
-    'backend', [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
+    "backend", [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
 )
 def test_backprop_glm_mc(laplace, model, reg_loader, backend):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True, backend=backend)
+    lap = laplace(model, "regression", enable_backprop=True, backend=backend)
     lap.fit(reg_loader)
-    f_mu, f_var = lap(X, pred_type='glm', link_approx='mc')
+    f_mu, f_var = lap(X, pred_type="glm", link_approx="mc")
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
@@ -725,17 +727,17 @@ def test_backprop_glm_mc(laplace, model, reg_loader, backend):
         assert False
 
 
-@pytest.mark.parametrize('laplace', [FullLaplace, KronLaplace, DiagLaplace])
+@pytest.mark.parametrize("laplace", [FullLaplace, KronLaplace, DiagLaplace])
 @pytest.mark.parametrize(
-    'backend', [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
+    "backend", [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
 )
 def test_backprop_nn(laplace, model, reg_loader, backend):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True, backend=backend)
+    lap = laplace(model, "regression", enable_backprop=True, backend=backend)
     lap.fit(reg_loader)
-    f_mu, f_var = lap(X, pred_type='nn', link_approx='mc', n_samples=10)
+    f_mu, f_var = lap(X, pred_type="nn", link_approx="mc", n_samples=10)
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
diff --git a/tests/test_curv_backends_asdfghjkl.py b/tests/test_curv_backends_asdfghjkl.py
index c22ce30..3d928bd 100644
--- a/tests/test_curv_backends_asdfghjkl.py
+++ b/tests/test_curv_backends_asdfghjkl.py
@@ -1,21 +1,20 @@
 import pytest
 import torch
+from asdfghjkl.operations import Bias, Scale
 from torch import nn
 from torch.nn.utils import parameters_to_vector
 
-from asdfghjkl.operations import Bias, Scale
-
-from laplace.curvature import AsdfghjklGGN, AsdfghjklEF, BackPackGGN, BackPackEF
+from laplace.curvature import AsdfghjklEF, AsdfghjklGGN, BackPackEF, BackPackGGN
 
 
 @pytest.fixture
 def model():
     torch.manual_seed(711)
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Tanh(), nn.Linear(20, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -40,10 +39,10 @@ def complex_model():
         Bias(),
         nn.Linear(20, 2),
     )
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -57,7 +56,7 @@ def complex_class_Xy():
 
 def test_diag_ggn_cls_kazuki_against_backpack_full(class_Xy, model):
     X, y = class_Xy
-    backend = AsdfghjklGGN(model, 'classification', stochastic=False)
+    backend = AsdfghjklGGN(model, "classification", stochastic=False)
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -67,7 +66,7 @@ def test_diag_ggn_cls_kazuki_against_backpack_full(class_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diagonal())
@@ -75,7 +74,7 @@ def test_diag_ggn_cls_kazuki_against_backpack_full(class_Xy, model):
 
 def test_diag_ef_cls_kazuki_against_backpack_full(class_Xy, model):
     X, y = class_Xy
-    backend = AsdfghjklEF(model, 'classification')
+    backend = AsdfghjklEF(model, "classification")
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -85,7 +84,7 @@ def test_diag_ef_cls_kazuki_against_backpack_full(class_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackEF(model, 'classification')
+    backend = BackPackEF(model, "classification")
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diagonal())
@@ -93,23 +92,23 @@ def test_diag_ef_cls_kazuki_against_backpack_full(class_Xy, model):
 
 def test_diag_ggn_stoch_cls_kazuki(class_Xy, model):
     X, y = class_Xy
-    backend = AsdfghjklGGN(model, 'classification', stochastic=True)
+    backend = AsdfghjklGGN(model, "classification", stochastic=True)
     loss, dggn = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
 
     # same order of magnitude os non-stochastic.
-    backend = AsdfghjklGGN(model, 'classification', stochastic=False)
+    backend = AsdfghjklGGN(model, "classification", stochastic=False)
     loss_ns, dggn_ns = backend.diag(X, y)
     assert loss_ns == loss
     assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1e1)
 
 
-@pytest.mark.parametrize('Backend', [AsdfghjklEF, AsdfghjklGGN])
+@pytest.mark.parametrize("Backend", [AsdfghjklEF, AsdfghjklGGN])
 def test_kron_kazuki_vs_diag_class(class_Xy, model, Backend):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = class_Xy
-    backend = Backend(model, 'classification')
+    backend = Backend(model, "classification")
     loss, dggn = backend.diag(X[:1], y[:1], N=1)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
@@ -117,10 +116,10 @@ def test_kron_kazuki_vs_diag_class(class_Xy, model, Backend):
     assert torch.allclose(kron.diag(), dggn)
 
 
-@pytest.mark.parametrize('Backend', [AsdfghjklEF, AsdfghjklGGN])
+@pytest.mark.parametrize("Backend", [AsdfghjklEF, AsdfghjklGGN])
 def test_kron_batching_correction_kazuki(class_Xy, model, Backend):
     X, y = class_Xy
-    backend = Backend(model, 'classification')
+    backend = Backend(model, "classification")
     loss, kron = backend.kron(X, y, N=len(X))
     assert len(kron.diag()) == model.n_params
 
@@ -134,11 +133,11 @@ def test_kron_batching_correction_kazuki(class_Xy, model, Backend):
     assert torch.allclose(loss, loss_two)
 
 
-@pytest.mark.parametrize('Backend', [AsdfghjklGGN, AsdfghjklEF])
+@pytest.mark.parametrize("Backend", [AsdfghjklGGN, AsdfghjklEF])
 def test_kron_summing_up_vs_diag_kazuki(class_Xy, model, Backend):
     # For a single data point, Kron is exact and should equal diag class_Xy
     X, y = class_Xy
-    backend = Backend(model, 'classification')
+    backend = Backend(model, "classification")
     loss, dggn = backend.diag(X, y, N=len(X))
     loss, kron = backend.kron(X, y, N=len(X))
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
@@ -146,7 +145,7 @@ def test_kron_summing_up_vs_diag_kazuki(class_Xy, model, Backend):
 
 def test_complex_diag_ggn_stoch_cls_kazuki(complex_class_Xy, complex_model):
     X, y = complex_class_Xy
-    backend = AsdfghjklGGN(complex_model, 'classification', stochastic=True)
+    backend = AsdfghjklGGN(complex_model, "classification", stochastic=True)
     loss, dggn = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(dggn) == complex_model.n_params
@@ -157,11 +156,11 @@ def test_complex_diag_ggn_stoch_cls_kazuki(complex_class_Xy, complex_model):
     assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1)
 
 
-@pytest.mark.parametrize('Backend', [AsdfghjklEF, AsdfghjklGGN])
+@pytest.mark.parametrize("Backend", [AsdfghjklEF, AsdfghjklGGN])
 def test_complex_kron_kazuki_vs_diag_kazuki(complex_class_Xy, complex_model, Backend):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = complex_class_Xy
-    backend = Backend(complex_model, 'classification')
+    backend = Backend(complex_model, "classification")
     loss, dggn = backend.diag(X[:1], y[:1], N=1)
     # sanity check size of diag ggn
     assert len(dggn) == complex_model.n_params
@@ -169,12 +168,12 @@ def test_complex_kron_kazuki_vs_diag_kazuki(complex_class_Xy, complex_model, Bac
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
 
 
-@pytest.mark.parametrize('Backend', [AsdfghjklEF, AsdfghjklGGN])
+@pytest.mark.parametrize("Backend", [AsdfghjklEF, AsdfghjklGGN])
 def test_complex_kron_batching_correction_kazuki(
     complex_class_Xy, complex_model, Backend
 ):
     X, y = complex_class_Xy
-    backend = Backend(complex_model, 'classification')
+    backend = Backend(complex_model, "classification")
     loss, kron = backend.kron(X, y, N=len(X))
     assert len(kron.diag()) == complex_model.n_params
 
@@ -188,13 +187,13 @@ def test_complex_kron_batching_correction_kazuki(
     assert torch.allclose(loss, loss_two)
 
 
-@pytest.mark.parametrize('Backend', [AsdfghjklGGN, AsdfghjklEF])
+@pytest.mark.parametrize("Backend", [AsdfghjklGGN, AsdfghjklEF])
 def test_complex_kron_summing_up_vs_diag_class_kazuki(
     complex_class_Xy, complex_model, Backend
 ):
     # For a single data point, Kron is exact and should equal diag class_Xy
     X, y = complex_class_Xy
-    backend = Backend(complex_model, 'classification')
+    backend = Backend(complex_model, "classification")
     loss, dggn = backend.diag(X, y, N=len(X))
     loss, kron = backend.kron(X, y, N=len(X))
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-2)
@@ -203,7 +202,7 @@ def test_complex_kron_summing_up_vs_diag_class_kazuki(
 def test_kron_normalization_ggn_class(class_Xy, model):
     X, y = class_Xy
     xi, yi = X[:1], y[:1]
-    backend = AsdfghjklGGN(model, 'classification', stochastic=False)
+    backend = AsdfghjklGGN(model, "classification", stochastic=False)
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
@@ -217,7 +216,7 @@ def test_kron_normalization_ggn_class(class_Xy, model):
 def test_kron_normalization_ef_class(class_Xy, model):
     X, y = class_Xy
     xi, yi = X[:1], y[:1]
-    backend = AsdfghjklEF(model, 'classification')
+    backend = AsdfghjklEF(model, "classification")
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
diff --git a/tests/test_curv_backends_asdl.py b/tests/test_curv_backends_asdl.py
index 4c9d7fa..e819e72 100644
--- a/tests/test_curv_backends_asdl.py
+++ b/tests/test_curv_backends_asdl.py
@@ -1,20 +1,21 @@
+from copy import deepcopy
+
 import pytest
 import torch
 from torch import nn
 from torch.nn.utils import parameters_to_vector
-from copy import deepcopy
 
-from laplace.curvature import AsdlGGN, AsdlEF, BackPackGGN, BackPackEF
+from laplace.curvature import AsdlEF, AsdlGGN, BackPackEF, BackPackGGN
 
 
 @pytest.fixture
 def model():
     torch.manual_seed(711)
     model = torch.nn.Sequential(nn.Linear(3, 3), nn.Tanh(), nn.Linear(3, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -45,10 +46,10 @@ def complex_model():
         nn.Tanh(),
         nn.Linear(20, 2),
     )
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -63,7 +64,7 @@ def complex_class_Xy():
 def test_diag_ggn_cls_asdl_against_backpack_full(class_Xy, model):
     X, y = class_Xy
     model2 = deepcopy(model)
-    backend = AsdlGGN(model, 'classification', stochastic=False)
+    backend = AsdlGGN(model, "classification", stochastic=False)
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -73,7 +74,7 @@ def test_diag_ggn_cls_asdl_against_backpack_full(class_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model2, 'classification', stochastic=False)
+    backend = BackPackGGN(model2, "classification", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diagonal())
@@ -82,7 +83,7 @@ def test_diag_ggn_cls_asdl_against_backpack_full(class_Xy, model):
 def test_diag_ggn_reg_asdl_against_backpack_full(reg_Xy, model):
     X, y = reg_Xy
     model2 = deepcopy(model)
-    backend = AsdlGGN(model, 'regression', stochastic=False)
+    backend = AsdlGGN(model, "regression", stochastic=False)
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -92,7 +93,7 @@ def test_diag_ggn_reg_asdl_against_backpack_full(reg_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model2, 'regression', stochastic=False)
+    backend = BackPackGGN(model2, "regression", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diagonal())
@@ -101,7 +102,7 @@ def test_diag_ggn_reg_asdl_against_backpack_full(reg_Xy, model):
 def test_diag_ef_cls_asdl_against_backpack_full(class_Xy, model):
     X, y = class_Xy
     model2 = deepcopy(model)
-    backend = AsdlEF(model, 'classification')
+    backend = AsdlEF(model, "classification")
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -111,7 +112,7 @@ def test_diag_ef_cls_asdl_against_backpack_full(class_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackEF(model2, 'classification')
+    backend = BackPackEF(model2, "classification")
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diagonal())
@@ -120,7 +121,7 @@ def test_diag_ef_cls_asdl_against_backpack_full(class_Xy, model):
 def test_diag_ef_reg_asdl_against_backpack_full(reg_Xy, model):
     X, y = reg_Xy
     model2 = deepcopy(model)
-    backend = AsdlEF(model, 'regression')
+    backend = AsdlEF(model, "regression")
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -130,7 +131,7 @@ def test_diag_ef_reg_asdl_against_backpack_full(reg_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackEF(model2, 'regression')
+    backend = BackPackEF(model2, "regression")
     loss_bp, dggn_bp = backend.diag(X, y)
     assert torch.allclose(loss, loss_bp)
     assert torch.allclose(dggn, dggn_bp)  # H_ggn.diagonal())
@@ -138,13 +139,13 @@ def test_diag_ef_reg_asdl_against_backpack_full(reg_Xy, model):
 
 def test_diag_ggn_stoch_cls_asdl(class_Xy, model):
     X, y = class_Xy
-    backend = AsdlGGN(model, 'classification', stochastic=True)
+    backend = AsdlGGN(model, "classification", stochastic=True)
     loss, dggn = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
 
     # same order of magnitude os non-stochastic.
-    backend = AsdlGGN(model, 'classification', stochastic=False)
+    backend = AsdlGGN(model, "classification", stochastic=False)
     loss_ns, dggn_ns = backend.diag(X, y)
     assert loss_ns == loss
     assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1e1)
@@ -153,12 +154,12 @@ def test_diag_ggn_stoch_cls_asdl(class_Xy, model):
 def test_kron_ggn_cls_asdl_against_backpack(class_Xy, model):
     X, y = class_Xy
     model2 = deepcopy(model)
-    backend = AsdlGGN(model, 'classification', stochastic=False)
+    backend = AsdlGGN(model, "classification", stochastic=False)
     loss, kron = backend.kron(X[:5], y[:5], len(y))
     loss2, kron2 = backend.kron(X[5:], y[5:], len(y))
     loss += loss2
     kron += kron2
-    backend_bp = BackPackGGN(model2, 'classification', stochastic=False)
+    backend_bp = BackPackGGN(model2, "classification", stochastic=False)
     loss_bp, kron_bp = backend_bp.kron(X, y, len(y))
 
     # sanity check size of diag ggn
@@ -170,13 +171,13 @@ def test_kron_ggn_cls_asdl_against_backpack(class_Xy, model):
 def test_kron_ggn_reg_asdl_against_backpack(reg_Xy, model):
     X, y = reg_Xy
     model2 = deepcopy(model)
-    backend = AsdlGGN(model, 'regression', stochastic=False)
+    backend = AsdlGGN(model, "regression", stochastic=False)
     # loss, kron = backend.kron(X, y, len(y))
     loss, kron = backend.kron(X[:5], y[:5], len(y))
     loss2, kron2 = backend.kron(X[5:], y[5:], len(y))
     loss += loss2
     kron += kron2
-    backend_bp = BackPackGGN(model2, 'regression', stochastic=False)
+    backend_bp = BackPackGGN(model2, "regression", stochastic=False)
     loss_bp, kron_bp = backend_bp.kron(X, y, len(y))
 
     # sanity check size of diag ggn
@@ -185,11 +186,11 @@ def test_kron_ggn_reg_asdl_against_backpack(reg_Xy, model):
     assert torch.allclose(kron.diag(), kron_bp.diag())
 
 
-@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
+@pytest.mark.parametrize("Backend", [AsdlEF, AsdlGGN])
 def test_kron_asdl_vs_diag_class(class_Xy, model, Backend):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = class_Xy
-    backend = Backend(model, 'classification')
+    backend = Backend(model, "classification")
     loss, dggn = backend.diag(X[:1], y[:1], N=1)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
@@ -197,11 +198,11 @@ def test_kron_asdl_vs_diag_class(class_Xy, model, Backend):
     assert torch.allclose(kron.diag(), dggn)
 
 
-@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
+@pytest.mark.parametrize("Backend", [AsdlEF, AsdlGGN])
 def test_kron_asdl_vs_diag_reg(reg_Xy, model, Backend):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = reg_Xy
-    backend = Backend(model, 'regression')
+    backend = Backend(model, "regression")
     loss, dggn = backend.diag(X[:1], y[:1], N=1)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
@@ -209,10 +210,10 @@ def test_kron_asdl_vs_diag_reg(reg_Xy, model, Backend):
     assert torch.allclose(kron.diag(), dggn)
 
 
-@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
+@pytest.mark.parametrize("Backend", [AsdlEF, AsdlGGN])
 def test_kron_batching_correction_asdl(class_Xy, model, Backend):
     X, y = class_Xy
-    backend = Backend(model, 'classification')
+    backend = Backend(model, "classification")
     loss, kron = backend.kron(X, y, N=len(X))
     assert len(kron.diag()) == model.n_params
 
@@ -226,11 +227,11 @@ def test_kron_batching_correction_asdl(class_Xy, model, Backend):
     assert torch.allclose(loss, loss_two)
 
 
-@pytest.mark.parametrize('Backend', [AsdlGGN, AsdlEF])
+@pytest.mark.parametrize("Backend", [AsdlGGN, AsdlEF])
 def test_kron_summing_up_vs_diag_asdl(class_Xy, model, Backend):
     # For a single data point, Kron is exact and should equal diag class_Xy
     X, y = class_Xy
-    backend = Backend(model, 'classification')
+    backend = Backend(model, "classification")
     loss, dggn = backend.diag(X, y, N=len(X))
     loss, kron = backend.kron(X, y, N=len(X))
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
@@ -238,7 +239,7 @@ def test_kron_summing_up_vs_diag_asdl(class_Xy, model, Backend):
 
 def test_complex_diag_ggn_stoch_cls_asdl(complex_class_Xy, complex_model):
     X, y = complex_class_Xy
-    backend = AsdlGGN(complex_model, 'classification', stochastic=True)
+    backend = AsdlGGN(complex_model, "classification", stochastic=True)
     loss, dggn = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(dggn) == complex_model.n_params
@@ -249,11 +250,11 @@ def test_complex_diag_ggn_stoch_cls_asdl(complex_class_Xy, complex_model):
     assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1)
 
 
-@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
+@pytest.mark.parametrize("Backend", [AsdlEF, AsdlGGN])
 def test_complex_kron_asdl_vs_diag_asdl(complex_class_Xy, complex_model, Backend):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = complex_class_Xy
-    backend = Backend(complex_model, 'classification')
+    backend = Backend(complex_model, "classification")
     loss, dggn = backend.diag(X[:1], y[:1], N=1)
     # sanity check size of diag ggn
     assert len(dggn) == complex_model.n_params
@@ -261,12 +262,12 @@ def test_complex_kron_asdl_vs_diag_asdl(complex_class_Xy, complex_model, Backend
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
 
 
-@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
+@pytest.mark.parametrize("Backend", [AsdlEF, AsdlGGN])
 def test_complex_kron_batching_correction_asdl(
     complex_class_Xy, complex_model, Backend
 ):
     X, y = complex_class_Xy
-    backend = Backend(complex_model, 'classification')
+    backend = Backend(complex_model, "classification")
     loss, kron = backend.kron(X, y, N=len(X))
     assert len(kron.diag()) == complex_model.n_params
 
@@ -280,13 +281,13 @@ def test_complex_kron_batching_correction_asdl(
     assert torch.allclose(loss, loss_two)
 
 
-@pytest.mark.parametrize('Backend', [AsdlGGN, AsdlEF])
+@pytest.mark.parametrize("Backend", [AsdlGGN, AsdlEF])
 def test_complex_kron_summing_up_vs_diag_class_asdl(
     complex_class_Xy, complex_model, Backend
 ):
     # For a single data point, Kron is exact and should equal diag class_Xy
     X, y = complex_class_Xy
-    backend = Backend(complex_model, 'classification')
+    backend = Backend(complex_model, "classification")
     loss, dggn = backend.diag(X, y, N=len(X))
     loss, kron = backend.kron(X, y, N=len(X))
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-2)
@@ -296,7 +297,7 @@ def test_kron_summation_ggn_class(class_Xy, model):
     X, y = class_Xy
     X = torch.repeat_interleave(X[:1], 7, 0)
     y = torch.repeat_interleave(y[:1], 7, 0)
-    backend = AsdlGGN(model, 'classification', stochastic=False)
+    backend = AsdlGGN(model, "classification", stochastic=False)
     loss_dg, diag = backend.diag(X, y)
     loss, kron = backend.kron(X, y, N=7)
     assert torch.allclose(loss_dg, loss)
@@ -307,7 +308,7 @@ def test_kron_summation_ggn_reg(reg_Xy, model):
     X, y = reg_Xy
     X = torch.repeat_interleave(X[:1], 7, 0)
     y = torch.repeat_interleave(y[:1], 7, 0)
-    backend = AsdlGGN(model, 'regression', stochastic=False)
+    backend = AsdlGGN(model, "regression", stochastic=False)
     loss_dg, diag = backend.diag(X, y)
     loss, kron = backend.kron(X, y, N=7)
     assert torch.allclose(loss_dg, loss)
@@ -317,7 +318,7 @@ def test_kron_summation_ggn_reg(reg_Xy, model):
 def test_kron_normalization_ggn_class(class_Xy, model):
     X, y = class_Xy
     xi, yi = X[:1], y[:1]
-    backend = AsdlGGN(model, 'classification', stochastic=False)
+    backend = AsdlGGN(model, "classification", stochastic=False)
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
@@ -331,7 +332,7 @@ def test_kron_normalization_ggn_class(class_Xy, model):
 def test_kron_normalization_ef_class(class_Xy, model):
     X, y = class_Xy
     xi, yi = X[:1], y[:1]
-    backend = AsdlEF(model, 'classification')
+    backend = AsdlEF(model, "classification")
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
@@ -345,7 +346,7 @@ def test_kron_normalization_ef_class(class_Xy, model):
 def test_kron_normalization_ggn_reg(reg_Xy, model):
     X, y = reg_Xy
     xi, yi = X[:1], y[:1]
-    backend = AsdlGGN(model, 'regression', stochastic=False)
+    backend = AsdlGGN(model, "regression", stochastic=False)
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
diff --git a/tests/test_curv_backends_backpack.py b/tests/test_curv_backends_backpack.py
index cd254a1..4dc54a9 100644
--- a/tests/test_curv_backends_backpack.py
+++ b/tests/test_curv_backends_backpack.py
@@ -3,17 +3,17 @@
 from torch import nn
 from torch.nn.utils import parameters_to_vector
 
-from laplace.curvature import BackPackGGN, BackPackEF
+from laplace.curvature import BackPackEF, BackPackGGN
 
 
 @pytest.fixture
 def model():
     torch.manual_seed(711)
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Tanh(), nn.Linear(20, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -21,10 +21,10 @@ def model():
 def model_singleoutput():
     torch.manual_seed(711)
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Tanh(), nn.Linear(20, 1))
-    setattr(model, 'output_size', 1)
+    setattr(model, "output_size", 1)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -48,7 +48,7 @@ def test_full_ggn_backpack_reg_integration(reg_Xy, model):
     X, y = reg_Xy
 
     # cannot test, its implemented based on Jacobians.
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss, H_ggn = backend.full(X, y)
     assert H_ggn.size() == torch.Size((model.n_params, model.n_params))
 
@@ -56,20 +56,20 @@ def test_full_ggn_backpack_reg_integration(reg_Xy, model):
 def test_full_ggn_backpack_cls_integration(class_Xy, model):
     X, y = class_Xy
     # cannot test, its implemented based on Jacobians.
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss, H_ggn = backend.full(X, y)
     assert H_ggn.size() == torch.Size((model.n_params, model.n_params))
 
 
 def test_diag_ggn_cls_backpack(class_Xy, model):
     X, y = class_Xy
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss, dggn = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert loss == loss_f
     assert torch.allclose(dggn, H_ggn.diagonal())
@@ -77,13 +77,13 @@ def test_diag_ggn_cls_backpack(class_Xy, model):
 
 def test_diag_ggn_reg_backpack(reg_Xy, model):
     X, y = reg_Xy
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss, dggn = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert loss == loss_f
     assert torch.allclose(dggn, H_ggn.diagonal())
@@ -91,13 +91,13 @@ def test_diag_ggn_reg_backpack(reg_Xy, model):
 
 def test_diag_ggn_stoch_cls_backpack(class_Xy, model):
     X, y = class_Xy
-    backend = BackPackGGN(model, 'classification', stochastic=True)
+    backend = BackPackGGN(model, "classification", stochastic=True)
     loss, dggn = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
 
     # same order of magnitude os non-stochastic.
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss_ns, dggn_ns = backend.diag(X, y)
     assert loss_ns == loss
     assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1e1)
@@ -105,13 +105,13 @@ def test_diag_ggn_stoch_cls_backpack(class_Xy, model):
 
 def test_diag_ggn_stoch_reg_backpack(reg_Xy, model):
     X, y = reg_Xy
-    backend = BackPackGGN(model, 'regression', stochastic=True)
+    backend = BackPackGGN(model, "regression", stochastic=True)
     loss, dggn = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
 
     # same order of magnitude os non-stochastic.
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss_ns, dggn_ns = backend.diag(X, y)
     assert loss_ns == loss
     assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1e1)
@@ -120,7 +120,7 @@ def test_diag_ggn_stoch_reg_backpack(reg_Xy, model):
 def test_kron_ggn_reg_backpack_vs_diag_reg(reg_Xy, model):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = reg_Xy
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss, dggn = backend.diag(X[:1], y[:1], N=1)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
@@ -130,7 +130,7 @@ def test_kron_ggn_reg_backpack_vs_diag_reg(reg_Xy, model):
 
 def test_kron_batching_correction_reg(reg_Xy, model):
     X, y = reg_Xy
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss, kron = backend.kron(X, y, N=len(X))
     assert len(kron.diag()) == model.n_params
 
@@ -147,7 +147,7 @@ def test_kron_batching_correction_reg(reg_Xy, model):
 def test_kron_summing_up_vs_diag_reg(reg_Xy, model):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = reg_Xy
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss, dggn = backend.diag(X, y, N=len(X))
     loss, kron = backend.kron(X, y, N=len(X))
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
@@ -156,7 +156,7 @@ def test_kron_summing_up_vs_diag_reg(reg_Xy, model):
 def test_kron_ggn_reg_backpack_vs_diag_class(class_Xy, model):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = class_Xy
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss, dggn = backend.diag(X[:1], y[:1], N=1)
     # sanity check size of diag ggn
     assert len(dggn) == model.n_params
@@ -166,7 +166,7 @@ def test_kron_ggn_reg_backpack_vs_diag_class(class_Xy, model):
 
 def test_kron_batching_correction_class(class_Xy, model):
     X, y = class_Xy
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss, kron = backend.kron(X, y, N=len(X))
     assert len(kron.diag()) == model.n_params
 
@@ -183,7 +183,7 @@ def test_kron_batching_correction_class(class_Xy, model):
 def test_kron_summing_up_vs_diag_class(class_Xy, model):
     # For a single data point, Kron is exact and should equal diag class_Xy
     X, y = class_Xy
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss, dggn = backend.diag(X, y, N=len(X))
     loss, kron = backend.kron(X, y, N=len(X))
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
@@ -191,13 +191,13 @@ def test_kron_summing_up_vs_diag_class(class_Xy, model):
 
 def test_full_vs_diag_ef_cls_backpack(class_Xy, model):
     X, y = class_Xy
-    backend = BackPackEF(model, 'classification')
+    backend = BackPackEF(model, "classification")
     loss, diag_ef = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(diag_ef) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackEF(model, 'classification')
+    backend = BackPackEF(model, "classification")
     loss_f, H_ef = backend.full(X, y)
     assert loss == loss_f
     assert torch.allclose(diag_ef, H_ef.diagonal())
@@ -205,13 +205,13 @@ def test_full_vs_diag_ef_cls_backpack(class_Xy, model):
 
 def test_full_vs_diag_ef_reg_backpack(reg_Xy, model):
     X, y = reg_Xy
-    backend = BackPackEF(model, 'regression')
+    backend = BackPackEF(model, "regression")
     loss, diag_ef = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(diag_ef) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackEF(model, 'regression')
+    backend = BackPackEF(model, "regression")
     loss_f, H_ef = backend.full(X, y)
     assert loss == loss_f
     assert torch.allclose(diag_ef, H_ef.diagonal())
@@ -220,13 +220,13 @@ def test_full_vs_diag_ef_reg_backpack(reg_Xy, model):
 def test_kron_normalization_reg(reg_Xy, model):
     X, y = reg_Xy
     xi, yi = X[:1], y[:1]
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
     X = torch.repeat_interleave(xi, 7, 0)
     y = torch.repeat_interleave(yi, 7, 0)
-    loss_test, kron_test  = backend.kron(X, y, N=7)
+    loss_test, kron_test = backend.kron(X, y, N=7)
     assert torch.allclose(kron_true.diag(), kron_test.diag())
     assert torch.allclose(loss_true, loss_test)
 
@@ -234,12 +234,12 @@ def test_kron_normalization_reg(reg_Xy, model):
 def test_kron_normalization_class(class_Xy, model):
     X, y = class_Xy
     xi, yi = X[:1], y[:1]
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
     X = torch.repeat_interleave(xi, 7, 0)
     y = torch.repeat_interleave(yi, 7, 0)
-    loss_test, kron_test  = backend.kron(X, y, N=7)
+    loss_test, kron_test = backend.kron(X, y, N=7)
     assert torch.allclose(kron_true.diag(), kron_test.diag())
     assert torch.allclose(loss_true, loss_test)
diff --git a/tests/test_curv_backends_curvlinops.py b/tests/test_curv_backends_curvlinops.py
index 2baaf10..0be1224 100644
--- a/tests/test_curv_backends_curvlinops.py
+++ b/tests/test_curv_backends_curvlinops.py
@@ -4,13 +4,13 @@
 from torch.nn.utils import parameters_to_vector
 
 from laplace.curvature import (
-    CurvlinopsGGN,
-    CurvlinopsEF,
-    CurvlinopsHessian,
-    BackPackGGN,
-    AsdlHessian,
     AsdlEF,
     AsdlGGN,
+    AsdlHessian,
+    BackPackGGN,
+    CurvlinopsEF,
+    CurvlinopsGGN,
+    CurvlinopsHessian,
 )
 
 
@@ -24,10 +24,10 @@ def run_around_tests():
 def model():
     torch.manual_seed(711)
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Tanh(), nn.Linear(20, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -50,10 +50,10 @@ def complex_model():
         nn.Tanh(),
         nn.Linear(20, 2),
     )
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -81,9 +81,9 @@ def complex_reg_Xy():
     return X, y
 
 
-@pytest.mark.parametrize('loss_type', ['classification', 'regression'])
+@pytest.mark.parametrize("loss_type", ["classification", "regression"])
 def test_full_hess_curvlinops_vs_asdl(class_Xy, reg_Xy, model, loss_type):
-    X, y = class_Xy if loss_type == 'classification' else reg_Xy
+    X, y = class_Xy if loss_type == "classification" else reg_Xy
 
     backend = CurvlinopsHessian(model, loss_type)
     loss, H = backend.full(X, y)
@@ -98,10 +98,10 @@ def test_full_hess_curvlinops_vs_asdl(class_Xy, reg_Xy, model, loss_type):
 def test_full_ggn_curvlinops_vs_asdl(class_Xy, model):
     X, y = class_Xy
 
-    backend = CurvlinopsGGN(model, 'classification', stochastic=False)
+    backend = CurvlinopsGGN(model, "classification", stochastic=False)
     loss, H = backend.full(X, y)
 
-    backend = AsdlGGN(model, 'classification', stochastic=False)
+    backend = AsdlGGN(model, "classification", stochastic=False)
     loss_ref, H_ref = backend.full(X, y)
 
     assert torch.allclose(loss, loss_ref, rtol=1e-4)
@@ -112,13 +112,13 @@ def test_full_ggn_stochastic(class_Xy, model):
     torch.manual_seed(123)
     X, y = class_Xy
 
-    backend = CurvlinopsGGN(model, 'classification', stochastic=True)
+    backend = CurvlinopsGGN(model, "classification", stochastic=True)
     loss_mc1, H_mc1 = backend.full(X, y, mc_samples=1)
 
-    backend = CurvlinopsGGN(model, 'classification', stochastic=True)
+    backend = CurvlinopsGGN(model, "classification", stochastic=True)
     loss_mc100, H_mc100 = backend.full(X, y, mc_samples=100)
 
-    backend = CurvlinopsGGN(model, 'classification', stochastic=False)
+    backend = CurvlinopsGGN(model, "classification", stochastic=False)
     loss_exact, H_exact = backend.full(X, y)
 
     assert torch.allclose(loss_mc1, loss_exact)
@@ -131,19 +131,19 @@ def test_full_ggn_stochastic(class_Xy, model):
 def test_full_ef_curvlinops_vs_asdl(class_Xy, model):
     X, y = class_Xy
 
-    backend = CurvlinopsEF(model, 'classification')
+    backend = CurvlinopsEF(model, "classification")
     loss, H = backend.full(X, y)
 
-    backend = AsdlEF(model, 'classification')
+    backend = AsdlEF(model, "classification")
     loss_ref, H_ref = backend.full(X, y)
 
     assert torch.allclose(loss, loss_ref, rtol=1e-4)
     assert torch.allclose(H, H_ref, rtol=1e-4)
 
 
-@pytest.mark.parametrize('loss_type', ['classification', 'regression'])
+@pytest.mark.parametrize("loss_type", ["classification", "regression"])
 def test_kron_ggn_curvlinops_vs_backpack(class_Xy, reg_Xy, model, loss_type):
-    X, y = class_Xy if loss_type == 'classification' else reg_Xy
+    X, y = class_Xy if loss_type == "classification" else reg_Xy
 
     backend = CurvlinopsGGN(model, loss_type, stochastic=False)
     loss, kron = backend.kron(X, y, N=1)
@@ -155,9 +155,9 @@ def test_kron_ggn_curvlinops_vs_backpack(class_Xy, reg_Xy, model, loss_type):
     assert torch.allclose(kron.to_matrix(), kron_ref.to_matrix(), rtol=5e-5)
 
 
-@pytest.mark.parametrize('loss_type', ['classification', 'regression'])
+@pytest.mark.parametrize("loss_type", ["classification", "regression"])
 def test_kron_ggn_stochastic(class_Xy, reg_Xy, model, loss_type):
-    X, y = class_Xy if loss_type == 'classification' else reg_Xy
+    X, y = class_Xy if loss_type == "classification" else reg_Xy
 
     backend = CurvlinopsGGN(model, loss_type, stochastic=True)
     loss_mc1, kron_mc1 = backend.kron(X, y, N=1, mc_samples=1)
@@ -176,17 +176,17 @@ def test_kron_ggn_stochastic(class_Xy, reg_Xy, model, loss_type):
     assert torch.norm(diff_mc1) > torch.norm(diff_mc100)
 
 
-@pytest.mark.parametrize('loss_type', ['classification', 'regression'])
+@pytest.mark.parametrize("loss_type", ["classification", "regression"])
 def test_kron_ggn_set_kfac_approx(
     complex_class_Xy, complex_reg_Xy, complex_model, loss_type
 ):
-    X, y = complex_class_Xy if loss_type == 'classification' else complex_reg_Xy
+    X, y = complex_class_Xy if loss_type == "classification" else complex_reg_Xy
 
     backend = CurvlinopsGGN(complex_model, loss_type, stochastic=False)
-    loss_expand, kron_expand = backend.kron(X, y, N=1, kfac_approx='expand')
+    loss_expand, kron_expand = backend.kron(X, y, N=1, kfac_approx="expand")
 
     backend = CurvlinopsGGN(complex_model, loss_type, stochastic=False)
-    loss_reduce, kron_reduce = backend.kron(X, y, N=1, kfac_approx='reduce')
+    loss_reduce, kron_reduce = backend.kron(X, y, N=1, kfac_approx="reduce")
 
     assert torch.allclose(loss_expand, loss_reduce)
     assert not torch.allclose(kron_expand.to_matrix(), kron_reduce.to_matrix())
@@ -194,20 +194,20 @@ def test_kron_ggn_set_kfac_approx(
 
 def test_kron_ef_cls_curvlinops_vs_backpack(class_Xy, model):
     X, y = class_Xy
-    backend = CurvlinopsEF(model, 'classification')
+    backend = CurvlinopsEF(model, "classification")
     loss, kron = backend.kron(X, y, N=1)
 
-    backend = AsdlEF(model, 'classification')
+    backend = AsdlEF(model, "classification")
     loss_ref, kron_ref = backend.kron(X, y, N=1)
 
     assert torch.allclose(loss, loss_ref)
     assert torch.allclose(kron.to_matrix(), kron_ref.to_matrix())
 
 
-@pytest.mark.parametrize('Backend', [CurvlinopsEF, CurvlinopsGGN])
+@pytest.mark.parametrize("Backend", [CurvlinopsEF, CurvlinopsGGN])
 def test_kron_batching_correction_cls_curvlinops(class_Xy, model, Backend):
     X, y = class_Xy
-    backend = Backend(model, 'classification')
+    backend = Backend(model, "classification")
     loss, kron = backend.kron(X, y, N=len(X))
     assert len(kron.diag()) == model.n_params
 
@@ -221,10 +221,10 @@ def test_kron_batching_correction_cls_curvlinops(class_Xy, model, Backend):
     assert torch.allclose(loss, loss_two)
 
 
-@pytest.mark.parametrize('Backend', [CurvlinopsEF, CurvlinopsGGN])
+@pytest.mark.parametrize("Backend", [CurvlinopsEF, CurvlinopsGGN])
 def test_kron_batching_correction_reg_curvlinops(reg_Xy, model, Backend):
     X, y = reg_Xy
-    backend = Backend(model, 'regression')
+    backend = Backend(model, "regression")
     loss, kron = backend.kron(X, y, N=len(X))
     assert len(kron.diag()) == model.n_params
 
@@ -238,10 +238,10 @@ def test_kron_batching_correction_reg_curvlinops(reg_Xy, model, Backend):
     assert torch.allclose(loss, loss_two)
 
 
-@pytest.mark.parametrize('Backend', [CurvlinopsEF, CurvlinopsGGN])
+@pytest.mark.parametrize("Backend", [CurvlinopsEF, CurvlinopsGGN])
 def test_kron_summing_up_vs_diag_cls_curvlinops(class_Xy, model, Backend):
     X, y = class_Xy
-    backend = Backend(model, 'classification')
+    backend = Backend(model, "classification")
     loss, dggn = backend.diag(X, y, N=len(X))
     loss, kron = backend.kron(X, y, N=len(X))
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
@@ -249,7 +249,7 @@ def test_kron_summing_up_vs_diag_cls_curvlinops(class_Xy, model, Backend):
 
 def test_complex_diag_ggn_stoch_cls_curvlinops(complex_class_Xy, complex_model):
     X, y = complex_class_Xy
-    backend = CurvlinopsGGN(complex_model, 'classification', stochastic=True)
+    backend = CurvlinopsGGN(complex_model, "classification", stochastic=True)
     loss, dggn = backend.diag(X, y)
     # sanity check size of diag ggn
     assert len(dggn) == complex_model.n_params
@@ -260,13 +260,13 @@ def test_complex_diag_ggn_stoch_cls_curvlinops(complex_class_Xy, complex_model):
     assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1)
 
 
-@pytest.mark.parametrize('Backend', [CurvlinopsEF, CurvlinopsGGN])
+@pytest.mark.parametrize("Backend", [CurvlinopsEF, CurvlinopsGGN])
 def test_complex_kron_curvlinops_vs_diag_curvlinops(
     complex_class_Xy, complex_model, Backend
 ):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = complex_class_Xy
-    backend = Backend(complex_model, 'classification')
+    backend = Backend(complex_model, "classification")
     loss, dggn = backend.diag(X[:1], y[:1], N=1)
     # sanity check size of diag ggn
     assert len(dggn) == complex_model.n_params
@@ -274,12 +274,12 @@ def test_complex_kron_curvlinops_vs_diag_curvlinops(
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
 
 
-@pytest.mark.parametrize('Backend', [CurvlinopsEF, CurvlinopsGGN])
+@pytest.mark.parametrize("Backend", [CurvlinopsEF, CurvlinopsGGN])
 def test_complex_kron_batching_correction_curvlinops(
     complex_class_Xy, complex_model, Backend
 ):
     X, y = complex_class_Xy
-    backend = Backend(complex_model, 'classification')
+    backend = Backend(complex_model, "classification")
     loss, kron = backend.kron(X, y, N=len(X))
     assert len(kron.diag()) == complex_model.n_params
 
@@ -293,13 +293,13 @@ def test_complex_kron_batching_correction_curvlinops(
     assert torch.allclose(loss, loss_two)
 
 
-@pytest.mark.parametrize('Backend', [CurvlinopsEF, CurvlinopsGGN])
+@pytest.mark.parametrize("Backend", [CurvlinopsEF, CurvlinopsGGN])
 def test_complex_kron_summing_up_vs_diag_class_curvlinops(
     complex_class_Xy, complex_model, Backend
 ):
     # For a single data point, Kron is exact and should equal diag class_Xy
     X, y = complex_class_Xy
-    backend = Backend(complex_model, 'classification')
+    backend = Backend(complex_model, "classification")
     loss, dggn = backend.diag(X, y, N=len(X))
     loss, kron = backend.kron(X, y, N=len(X))
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-2)
@@ -308,7 +308,7 @@ def test_complex_kron_summing_up_vs_diag_class_curvlinops(
 def test_kron_normalization_ggn_class(class_Xy, model):
     X, y = class_Xy
     xi, yi = X[:1], y[:1]
-    backend = CurvlinopsGGN(model, 'classification', stochastic=False)
+    backend = CurvlinopsGGN(model, "classification", stochastic=False)
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
@@ -322,7 +322,7 @@ def test_kron_normalization_ggn_class(class_Xy, model):
 def test_kron_normalization_ef_class(class_Xy, model):
     X, y = class_Xy
     xi, yi = X[:1], y[:1]
-    backend = CurvlinopsEF(model, 'classification')
+    backend = CurvlinopsEF(model, "classification")
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
diff --git a/tests/test_curv_backends_interface.py b/tests/test_curv_backends_interface.py
index d3270b9..475d28b 100644
--- a/tests/test_curv_backends_interface.py
+++ b/tests/test_curv_backends_interface.py
@@ -1,21 +1,17 @@
-import pytest
 import numpy as np
+import pytest
 import torch
 from torch import nn
 from torch.nn.utils import parameters_to_vector
 
-from asdl.operations import Bias, Scale
-
 from laplace.curvature import (
-    CurvlinopsGGN,
-    CurvlinopsEF,
-    BackPackGGN,
-    BackPackEF,
-    GGNInterface,
-    EFInterface,
-    AsdlGGN,
     AsdlEF,
+    AsdlGGN,
+    BackPackEF,
+    BackPackGGN,
     CurvatureInterface,
+    EFInterface,
+    GGNInterface,
 )
 
 
@@ -23,10 +19,10 @@
 def model():
     torch.manual_seed(711)
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Tanh(), nn.Linear(20, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -57,10 +53,10 @@ def complex_model():
         nn.Tanh(),
         nn.Linear(20, 2),
     )
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -74,10 +70,10 @@ def complex_class_Xy():
 
 def test_batchgrad_cls(class_Xy, model):
     X, y = class_Xy
-    backend = CurvatureInterface(model, 'classification')
+    backend = CurvatureInterface(model, "classification")
     batchgrads, loss = backend.gradients(X, y)
 
-    backend_asdl = AsdlEF(model, 'classification')
+    backend_asdl = AsdlEF(model, "classification")
     batchgrads_ref, loss_ref = backend_asdl.gradients(X, y)
 
     torch.allclose(batchgrads, batchgrads_ref)
@@ -85,10 +81,10 @@ def test_batchgrad_cls(class_Xy, model):
 
 def test_batchgrad_cls_complex(complex_class_Xy, complex_model):
     X, y = complex_class_Xy
-    backend = CurvatureInterface(complex_model, 'classification')
+    backend = CurvatureInterface(complex_model, "classification")
     batchgrads, loss = backend.gradients(X, y)
 
-    backend_asdl = AsdlEF(complex_model, 'classification')
+    backend_asdl = AsdlEF(complex_model, "classification")
     batchgrads_ref, loss_ref = backend_asdl.gradients(X, y)
 
     torch.allclose(batchgrads, batchgrads_ref)
@@ -96,10 +92,10 @@ def test_batchgrad_cls_complex(complex_class_Xy, complex_model):
 
 def test_batchgrad_reg(reg_Xy, model):
     X, y = reg_Xy
-    backend = CurvatureInterface(model, 'regression')
+    backend = CurvatureInterface(model, "regression")
     batchgrads, loss = backend.gradients(X, y)
 
-    backend_asdl = BackPackEF(model, 'regression')
+    backend_asdl = BackPackEF(model, "regression")
     batchgrads_ref, loss_ref = backend_asdl.gradients(X, y)
 
     torch.allclose(batchgrads, batchgrads_ref)
@@ -109,7 +105,7 @@ def test_diag_ggn_cls_curvlinops_against_backpack_full(class_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = class_Xy
-    backend = GGNInterface(model, 'classification', stochastic=False)
+    backend = GGNInterface(model, "classification", stochastic=False)
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -119,7 +115,7 @@ def test_diag_ggn_cls_curvlinops_against_backpack_full(class_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diag())
@@ -131,7 +127,7 @@ def test_diag_ggn_complex_cls_curvlinops_against_backpack_full(
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = complex_class_Xy
-    backend = GGNInterface(complex_model, 'classification', stochastic=False)
+    backend = GGNInterface(complex_model, "classification", stochastic=False)
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -141,7 +137,7 @@ def test_diag_ggn_complex_cls_curvlinops_against_backpack_full(
     assert len(dggn) == complex_model.n_params
 
     # check against manually computed full GGN:
-    backend = AsdlGGN(complex_model, 'classification', stochastic=False)
+    backend = AsdlGGN(complex_model, "classification", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diag())
@@ -151,7 +147,7 @@ def test_diag_ggn_stoch_cls_curvlinops_against_backpack_full(class_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = class_Xy
-    backend = GGNInterface(model, 'classification', stochastic=True, num_samples=10000)
+    backend = GGNInterface(model, "classification", stochastic=True, num_samples=10000)
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -161,7 +157,7 @@ def test_diag_ggn_stoch_cls_curvlinops_against_backpack_full(class_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diag(), atol=0.01)
@@ -171,7 +167,7 @@ def test_diag_ef_cls_curvlinops_against_backpack_full(class_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = class_Xy
-    backend = EFInterface(model, 'classification')
+    backend = EFInterface(model, "classification")
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -181,7 +177,7 @@ def test_diag_ef_cls_curvlinops_against_backpack_full(class_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackEF(model, 'classification')
+    backend = BackPackEF(model, "classification")
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diag())
@@ -191,7 +187,7 @@ def test_diag_ggn_reg_curvlinops_against_backpack_full(reg_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = reg_Xy
-    backend = GGNInterface(model, 'regression', stochastic=False)
+    backend = GGNInterface(model, "regression", stochastic=False)
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -201,7 +197,7 @@ def test_diag_ggn_reg_curvlinops_against_backpack_full(reg_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diag())
@@ -211,7 +207,7 @@ def test_diag_ggn_stoch_reg_curvlinops_against_backpack_full(reg_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = reg_Xy
-    backend = GGNInterface(model, 'regression', stochastic=True, num_samples=10000)
+    backend = GGNInterface(model, "regression", stochastic=True, num_samples=10000)
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -221,7 +217,7 @@ def test_diag_ggn_stoch_reg_curvlinops_against_backpack_full(reg_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diag(), atol=0.1)
@@ -231,7 +227,7 @@ def test_diag_ef_reg_curvlinops_against_backpack_full(reg_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = reg_Xy
-    backend = EFInterface(model, 'regression')
+    backend = EFInterface(model, "regression")
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
     loss += loss2
@@ -241,7 +237,7 @@ def test_diag_ef_reg_curvlinops_against_backpack_full(reg_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackEF(model, 'regression')
+    backend = BackPackEF(model, "regression")
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diag())
@@ -251,7 +247,7 @@ def test_full_ggn_cls_curvlinops_against_backpack_full(class_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = class_Xy
-    backend = GGNInterface(model, 'classification', stochastic=False)
+    backend = GGNInterface(model, "classification", stochastic=False)
     loss, fggn = backend.full(X[:5], y[:5])
     loss2, fggn2 = backend.full(X[5:], y[5:])
     loss += loss2
@@ -261,7 +257,7 @@ def test_full_ggn_cls_curvlinops_against_backpack_full(class_Xy, model):
     assert fggn.shape == (model.n_params, model.n_params)
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(fggn, H_ggn, atol=0.1)
@@ -271,7 +267,7 @@ def test_full_ggn_stoch_cls_curvlinops_against_backpack_full(class_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = class_Xy
-    backend = GGNInterface(model, 'classification', stochastic=True, num_samples=10000)
+    backend = GGNInterface(model, "classification", stochastic=True, num_samples=10000)
     loss, fggn = backend.full(X[:5], y[:5])
     loss2, fggn2 = backend.full(X[5:], y[5:])
     loss += loss2
@@ -281,7 +277,7 @@ def test_full_ggn_stoch_cls_curvlinops_against_backpack_full(class_Xy, model):
     assert fggn.shape == (model.n_params, model.n_params)
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model, "classification", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(fggn, H_ggn, atol=0.1)
@@ -291,7 +287,7 @@ def test_full_ef_cls_curvlinops_against_backpack_full(class_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = class_Xy
-    backend = EFInterface(model, 'classification')
+    backend = EFInterface(model, "classification")
     loss, fggn = backend.full(X[:5], y[:5])
     loss2, fggn2 = backend.full(X[5:], y[5:])
     loss += loss2
@@ -301,7 +297,7 @@ def test_full_ef_cls_curvlinops_against_backpack_full(class_Xy, model):
     assert fggn.shape == (model.n_params, model.n_params)
 
     # check against manually computed full GGN:
-    backend = BackPackEF(model, 'classification')
+    backend = BackPackEF(model, "classification")
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(fggn, H_ggn, atol=0.0001)
@@ -311,7 +307,7 @@ def test_full_ggn_reg_curvlinops_against_backpack_full(reg_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = reg_Xy
-    backend = GGNInterface(model, 'regression', stochastic=False)
+    backend = GGNInterface(model, "regression", stochastic=False)
     loss, fggn = backend.full(X[:5], y[:5])
     loss2, fggn2 = backend.full(X[5:], y[5:])
     loss += loss2
@@ -321,7 +317,7 @@ def test_full_ggn_reg_curvlinops_against_backpack_full(reg_Xy, model):
     assert fggn.shape == (model.n_params, model.n_params)
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(fggn, H_ggn, atol=0.1)
@@ -331,7 +327,7 @@ def test_full_ggn_stoch_reg_curvlinops_against_backpack_full(reg_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = reg_Xy
-    backend = GGNInterface(model, 'regression', stochastic=True, num_samples=10000)
+    backend = GGNInterface(model, "regression", stochastic=True, num_samples=10000)
     loss, fggn = backend.full(X[:5], y[:5])
     loss2, fggn2 = backend.full(X[5:], y[5:])
     loss += loss2
@@ -341,7 +337,7 @@ def test_full_ggn_stoch_reg_curvlinops_against_backpack_full(reg_Xy, model):
     assert fggn.shape == (model.n_params, model.n_params)
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(fggn, H_ggn, atol=0.1)
@@ -351,7 +347,7 @@ def test_full_ef_reg_curvlinops_against_backpack_full(reg_Xy, model):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = reg_Xy
-    backend = EFInterface(model, 'regression')
+    backend = EFInterface(model, "regression")
     loss, fggn = backend.full(X[:5], y[:5])
     loss2, fggn2 = backend.full(X[5:], y[5:])
     loss += loss2
@@ -361,7 +357,7 @@ def test_full_ef_reg_curvlinops_against_backpack_full(reg_Xy, model):
     assert fggn.shape == (model.n_params, model.n_params)
 
     # check against manually computed full GGN:
-    backend = BackPackEF(model, 'regression')
+    backend = BackPackEF(model, "regression")
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(fggn, H_ggn, atol=0.0001)
diff --git a/tests/test_feature_extractor.py b/tests/test_feature_extractor.py
index d3c3086..6dca5f8 100644
--- a/tests/test_feature_extractor.py
+++ b/tests/test_feature_extractor.py
@@ -1,11 +1,9 @@
+import pytest
 import torch
 import torch.nn as nn
 import torchvision.models as models
 
 from laplace.utils import FeatureExtractor
-
-import pytest
-
 from laplace.utils.feature_extractor import FeatureReduction
 
 
@@ -52,37 +50,37 @@ def forward(self, x):
 
 
 def get_model(model_name):
-    if model_name == 'resnet18':
+    if model_name == "resnet18":
         model = models.resnet18()
-    elif model_name == 'alexnet':
+    elif model_name == "alexnet":
         model = models.alexnet()
-    elif model_name == 'vgg16':
+    elif model_name == "vgg16":
         model = models.vgg16()
-    elif model_name == 'squeezenet':
+    elif model_name == "squeezenet":
         model = models.squeezenet1_0()
-    elif model_name == 'densenet':
+    elif model_name == "densenet":
         model = models.densenet161()
-    elif model_name == 'inception':
+    elif model_name == "inception":
         model = models.inception_v3(init_weights=True)
-    elif model_name == 'googlenet':
+    elif model_name == "googlenet":
         model = models.googlenet(init_weights=True)
-    elif model_name == 'shufflenet':
+    elif model_name == "shufflenet":
         model = models.shufflenet_v2_x1_0()
-    elif model_name == 'mobilenet_v2':
+    elif model_name == "mobilenet_v2":
         model = models.mobilenet_v2()
-    elif model_name == 'mobilenet_v3_large':
+    elif model_name == "mobilenet_v3_large":
         model = models.mobilenet_v3_large()
-    elif model_name == 'mobilenet_v3_small':
+    elif model_name == "mobilenet_v3_small":
         model = models.mobilenet_v3_small()
-    elif model_name == 'resnext50_32x4d':
+    elif model_name == "resnext50_32x4d":
         model = models.resnext50_32x4d()
-    elif model_name == 'wide_resnet50_2':
+    elif model_name == "wide_resnet50_2":
         model = models.wide_resnet50_2()
-    elif model_name == 'mnasnet':
+    elif model_name == "mnasnet":
         model = models.mnasnet1_0()
-    elif model_name == 'switchedCNN':
+    elif model_name == "switchedCNN":
         model = CNN(10)
-    elif model_name == 'sequential':
+    elif model_name == "sequential":
         model = nn.Sequential(
             nn.Conv2d(3, 6, 3, 1, 1),
             nn.ReLU(),
@@ -92,7 +90,7 @@ def get_model(model_name):
             nn.Linear(10, 10),
         )
     else:
-        raise ValueError(f'{model_name} is not supported.')
+        raise ValueError(f"{model_name} is not supported.")
     return model
 
 
@@ -101,21 +99,21 @@ def test_feature_extractor():
     # all torchvision classifcation models but 'squeezenet' (no linear last layer)
     # + model where modules are initilaized in wrong order + nn.Sequential model
     model_names = [
-        'resnet18',
-        'alexnet',
-        'vgg16',
-        'densenet',
-        'inception',
-        'googlenet',
-        'shufflenet',
-        'mobilenet_v2',
-        'mobilenet_v3_large',
-        'mobilenet_v3_small',
-        'resnext50_32x4d',
-        'wide_resnet50_2',
-        'mnasnet',
-        'switchedCNN',
-        'sequential',
+        "resnet18",
+        "alexnet",
+        "vgg16",
+        "densenet",
+        "inception",
+        "googlenet",
+        "shufflenet",
+        "mobilenet_v2",
+        "mobilenet_v3_large",
+        "mobilenet_v3_small",
+        "resnext50_32x4d",
+        "wide_resnet50_2",
+        "mnasnet",
+        "switchedCNN",
+        "sequential",
     ]
 
     # to test the last_layer_name argument
@@ -125,7 +123,7 @@ def test_feature_extractor():
 
     for model_name in model_names:
         # generate random test input
-        if model_name == 'inception':
+        if model_name == "inception":
             x = torch.randn(1, 3, 299, 299)
         else:
             x = torch.randn(1, 3, 64, 64)
@@ -144,8 +142,8 @@ def test_feature_extractor():
 
 
 @torch.no_grad()
-@pytest.mark.parametrize('reduction', [r.value for r in FeatureReduction] + [None])
-@pytest.mark.parametrize('additional_dims', [tuple(), (7,), (7, 8, 9)])
+@pytest.mark.parametrize("reduction", [r.value for r in FeatureReduction] + [None])
+@pytest.mark.parametrize("additional_dims", [tuple(), (7,), (7, 8, 9)])
 def test_multidim_features(reduction, additional_dims):
     BATCH_SIZE = 6
     IN_DIM = 5
diff --git a/tests/test_jacobians.py b/tests/test_jacobians.py
index b32bfa4..0d08abc 100644
--- a/tests/test_jacobians.py
+++ b/tests/test_jacobians.py
@@ -10,21 +10,21 @@
 @pytest.fixture
 def multioutput_model():
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     return model
 
 
 @pytest.fixture
 def singleoutput_model():
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 1))
-    setattr(model, 'output_size', 1)
+    setattr(model, "output_size", 1)
     return model
 
 
 @pytest.fixture
 def linear_model():
     model = torch.nn.Sequential(nn.Linear(3, 1, bias=False))
-    setattr(model, 'output_size', 1)
+    setattr(model, "output_size", 1)
     return model
 
 
@@ -35,11 +35,11 @@ def X():
 
 
 @pytest.mark.parametrize(
-    'backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface]
+    "backend_cls", [CurvatureInterface, AsdlInterface, BackPackInterface]
 )
 def test_linear_jacobians(linear_model, X, backend_cls):
     # jacobian of linear model is input X.
-    backend = backend_cls(linear_model, 'classification')
+    backend = backend_cls(linear_model, "classification")
     Js, f = backend.jacobians(X)
     # into Jacs shape (batch_size, output_size, params)
     true_Js = X.reshape(len(X), 1, -1)
@@ -49,11 +49,11 @@ def test_linear_jacobians(linear_model, X, backend_cls):
 
 
 @pytest.mark.parametrize(
-    'backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface]
+    "backend_cls", [CurvatureInterface, AsdlInterface, BackPackInterface]
 )
 def test_jacobians_singleoutput(singleoutput_model, X, backend_cls):
     model = singleoutput_model
-    backend = backend_cls(model, 'classification')
+    backend = backend_cls(model, "classification")
     Js, f = backend.jacobians(X)
     Js_naive, f_naive = jacobians_naive(model, X)
     assert Js.shape == Js_naive.shape
@@ -63,11 +63,11 @@ def test_jacobians_singleoutput(singleoutput_model, X, backend_cls):
 
 
 @pytest.mark.parametrize(
-    'backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface]
+    "backend_cls", [CurvatureInterface, AsdlInterface, BackPackInterface]
 )
 def test_jacobians_multioutput(multioutput_model, X, backend_cls):
     model = multioutput_model
-    backend = backend_cls(model, 'classification')
+    backend = backend_cls(model, "classification")
     Js, f = backend.jacobians(X)
     Js_naive, f_naive = jacobians_naive(model, X)
     assert Js.shape == Js_naive.shape
@@ -77,11 +77,11 @@ def test_jacobians_multioutput(multioutput_model, X, backend_cls):
 
 
 @pytest.mark.parametrize(
-    'backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface]
+    "backend_cls", [CurvatureInterface, AsdlInterface, BackPackInterface]
 )
 def test_last_layer_jacobians_singleoutput(singleoutput_model, X, backend_cls):
     model = FeatureExtractor(singleoutput_model)
-    backend = backend_cls(model, 'classification')
+    backend = backend_cls(model, "classification")
     Js, f = backend.last_layer_jacobians(X)
     _, phi = model.forward_with_features(X)
     Js_naive, f_naive = jacobians_naive(model.last_layer, phi)
@@ -92,11 +92,11 @@ def test_last_layer_jacobians_singleoutput(singleoutput_model, X, backend_cls):
 
 
 @pytest.mark.parametrize(
-    'backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface]
+    "backend_cls", [CurvatureInterface, AsdlInterface, BackPackInterface]
 )
 def test_last_layer_jacobians_multioutput(multioutput_model, X, backend_cls):
     model = FeatureExtractor(multioutput_model)
-    backend = backend_cls(model, 'classification')
+    backend = backend_cls(model, "classification")
     Js, f = backend.last_layer_jacobians(X)
     _, phi = model.forward_with_features(X)
     Js_naive, f_naive = jacobians_naive(model.last_layer, phi)
@@ -106,10 +106,10 @@ def test_last_layer_jacobians_multioutput(multioutput_model, X, backend_cls):
     assert torch.allclose(f, f_naive)
 
 
-@pytest.mark.parametrize('backend_cls', [CurvatureInterface, BackPackInterface])
+@pytest.mark.parametrize("backend_cls", [CurvatureInterface, BackPackInterface])
 def test_backprop_jacobians_singleoutput(singleoutput_model, X, backend_cls):
     X.requires_grad = True
-    backend = backend_cls(singleoutput_model, 'regression')
+    backend = backend_cls(singleoutput_model, "regression")
 
     try:
         Js, f = backend.jacobians(X, enable_backprop=True)
@@ -124,10 +124,10 @@ def test_backprop_jacobians_singleoutput(singleoutput_model, X, backend_cls):
         assert False
 
 
-@pytest.mark.parametrize('backend_cls', [CurvatureInterface, BackPackInterface])
+@pytest.mark.parametrize("backend_cls", [CurvatureInterface, BackPackInterface])
 def test_backprop_jacobians_multioutput(multioutput_model, X, backend_cls):
     X.requires_grad = True
-    backend = backend_cls(multioutput_model, 'regression')
+    backend = backend_cls(multioutput_model, "regression")
 
     try:
         Js, f = backend.jacobians(X, enable_backprop=True)
@@ -143,11 +143,11 @@ def test_backprop_jacobians_multioutput(multioutput_model, X, backend_cls):
         assert False
 
 
-@pytest.mark.parametrize('backend_cls', [CurvatureInterface, BackPackInterface])
+@pytest.mark.parametrize("backend_cls", [CurvatureInterface, BackPackInterface])
 def test_backprop_last_layer_jacobians_singleoutput(singleoutput_model, X, backend_cls):
     X.requires_grad = True
     backend = backend_cls(
-        FeatureExtractor(singleoutput_model, enable_backprop=True), 'regression'
+        FeatureExtractor(singleoutput_model, enable_backprop=True), "regression"
     )
 
     try:
@@ -163,11 +163,11 @@ def test_backprop_last_layer_jacobians_singleoutput(singleoutput_model, X, backe
         assert False
 
 
-@pytest.mark.parametrize('backend_cls', [CurvatureInterface, BackPackInterface])
+@pytest.mark.parametrize("backend_cls", [CurvatureInterface, BackPackInterface])
 def test_backprop_last_layer_jacobians_multioutput(multioutput_model, X, backend_cls):
     X.requires_grad = True
     backend = backend_cls(
-        FeatureExtractor(multioutput_model, enable_backprop=True), 'regression'
+        FeatureExtractor(multioutput_model, enable_backprop=True), "regression"
     )
 
     try:
diff --git a/tests/test_laplace.py b/tests/test_laplace.py
index c40be16..1aed014 100644
--- a/tests/test_laplace.py
+++ b/tests/test_laplace.py
@@ -3,47 +3,64 @@
 from torch import nn
 from torch.nn.utils import parameters_to_vector
 
+from laplace.baselaplace import DiagLaplace, FullLaplace, KronLaplace
 from laplace.laplace import Laplace
-from laplace.baselaplace import FullLaplace, KronLaplace, DiagLaplace
-from laplace.lllaplace import FullLLLaplace, KronLLLaplace, DiagLLLaplace
-
+from laplace.lllaplace import DiagLLLaplace, FullLLLaplace, KronLLLaplace
 
 torch.manual_seed(240)
 torch.set_default_tensor_type(torch.DoubleTensor)
-flavors = [FullLaplace, KronLaplace, DiagLaplace,
-           FullLLLaplace, KronLLLaplace, DiagLLLaplace]
-all_keys = [('all', 'full'), ('all', 'kron'), ('all', 'diag'),
-            ('last_layer', 'full'), ('last_layer', 'kron'), ('last_layer', 'diag')]
+flavors = [
+    FullLaplace,
+    KronLaplace,
+    DiagLaplace,
+    FullLLLaplace,
+    KronLLLaplace,
+    DiagLLLaplace,
+]
+all_keys = [
+    ("all", "full"),
+    ("all", "kron"),
+    ("all", "diag"),
+    ("last_layer", "full"),
+    ("last_layer", "kron"),
+    ("last_layer", "diag"),
+]
 
 
 @pytest.fixture
 def model():
     model = nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
-    model_params = list(model.parameters())
     return model
 
 
-def test_default_init(model, likelihood='classification'):
+def test_default_init(model, likelihood="classification"):
     # test if default initialization works, id=(last-layer, kron)
     lap = Laplace(model, likelihood)
     assert isinstance(lap, KronLLLaplace)
 
 
-@pytest.mark.parametrize('laplace, key', zip(flavors, all_keys))
-def test_all_init(laplace, key, model, likelihood='classification'):
+@pytest.mark.parametrize("laplace, key", zip(flavors, all_keys))
+def test_all_init(laplace, key, model, likelihood="classification"):
     # test if all flavors are correctly initialized
     w, s = key
     lap = Laplace(model, likelihood, subset_of_weights=w, hessian_structure=s)
     assert isinstance(lap, laplace)
 
 
-@pytest.mark.parametrize('key', all_keys)
-def test_opt_keywords(key, model, likelihood='classification'):
+@pytest.mark.parametrize("key", all_keys)
+def test_opt_keywords(key, model, likelihood="classification"):
     # test if optional keywords are correctly passed on
     w, s = key
     prior_mean = torch.zeros_like(parameters_to_vector(model.parameters()))
-    lap = Laplace(model, likelihood, subset_of_weights=w, hessian_structure=s,
-                  prior_precision=0.01, prior_mean=prior_mean, temperature=10.)
+    lap = Laplace(
+        model,
+        likelihood,
+        subset_of_weights=w,
+        hessian_structure=s,
+        prior_precision=0.01,
+        prior_mean=prior_mean,
+        temperature=10.0,
+    )
     assert torch.allclose(lap.prior_mean, prior_mean)
     assert lap.prior_precision == 0.01
-    assert lap.temperature == 10.
+    assert lap.temperature == 10.0
diff --git a/tests/test_lllaplace.py b/tests/test_lllaplace.py
index d5f7d38..c1596f1 100644
--- a/tests/test_lllaplace.py
+++ b/tests/test_lllaplace.py
@@ -27,7 +27,7 @@ def run_around_tests():
 @pytest.fixture
 def model():
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     return model
 
 
@@ -83,9 +83,9 @@ def multidim_reg_loader():
     return DataLoader(TensorDataset(X, y), batch_size=3)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init(laplace, model):
-    lap = laplace(model, 'classification', last_layer_name='1')
+    lap = laplace(model, "classification", last_layer_name="1")
     assert torch.allclose(lap.mean, lap.prior_mean)
     if laplace != KronLLLaplace:
         H = lap.H.clone()
@@ -99,16 +99,16 @@ def test_laplace_init(laplace, model):
                 assert torch.allclose(k1, k2)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init_nollname(laplace, model):
-    lap = laplace(model, 'classification')
+    lap = laplace(model, "classification")
     assert lap.mean is None
     assert lap.H is None
 
 
-@pytest.mark.parametrize('laplace', [KronLLLaplace, DiagLLLaplace])
+@pytest.mark.parametrize("laplace", [KronLLLaplace, DiagLLLaplace])
 def test_laplace_large_init(laplace, large_model):
-    lap = laplace(large_model, 'classification', last_layer_name='fc')
+    lap = laplace(large_model, "classification", last_layer_name="fc")
     assert torch.allclose(lap.mean, lap.prior_mean)
     if laplace == DiagLLLaplace:
         H = lap.H.clone()
@@ -122,35 +122,35 @@ def test_laplace_large_init(laplace, large_model):
                 assert torch.allclose(k1, k2)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_large_init_nollname(laplace, large_model):
-    lap = laplace(large_model, 'classification')
+    lap = laplace(large_model, "classification")
     assert lap.mean is None
     assert lap.H is None
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_invalid_likelihood(laplace, model):
     with pytest.raises(ValueError):
-        laplace(model, 'otherlh', last_layer_name='1')
+        laplace(model, "otherlh", last_layer_name="1")
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init_noise(laplace, model):
     # float
     sigma_noise = 1.2
     laplace(
-        model, likelihood='regression', sigma_noise=sigma_noise, last_layer_name='1'
+        model, likelihood="regression", sigma_noise=sigma_noise, last_layer_name="1"
     )
     # torch.tensor 0-dim
     sigma_noise = torch.tensor(1.2)
     laplace(
-        model, likelihood='regression', sigma_noise=sigma_noise, last_layer_name='1'
+        model, likelihood="regression", sigma_noise=sigma_noise, last_layer_name="1"
     )
     # torch.tensor 1-dim
     sigma_noise = torch.tensor(1.2).reshape(-1)
     laplace(
-        model, likelihood='regression', sigma_noise=sigma_noise, last_layer_name='1'
+        model, likelihood="regression", sigma_noise=sigma_noise, last_layer_name="1"
     )
 
     # for classification should fail
@@ -158,9 +158,9 @@ def test_laplace_init_noise(laplace, model):
     with pytest.raises(ValueError):
         laplace(
             model,
-            likelihood='classification',
+            likelihood="classification",
             sigma_noise=sigma_noise,
-            last_layer_name='1',
+            last_layer_name="1",
         )
 
     # other than that should fail
@@ -168,121 +168,122 @@ def test_laplace_init_noise(laplace, model):
     sigma_noise = torch.tensor(1.2).reshape(1, 1)
     with pytest.raises(ValueError):
         laplace(
-            model, likelihood='regression', sigma_noise=sigma_noise, last_layer_name='1'
+            model, likelihood="regression", sigma_noise=sigma_noise, last_layer_name="1"
         )
     # other datatype, only reals supported
-    sigma_noise = '1.2'
+    sigma_noise = "1.2"
     with pytest.raises(ValueError):
         laplace(
-            model, likelihood='regression', sigma_noise=sigma_noise, last_layer_name='1'
+            model, likelihood="regression", sigma_noise=sigma_noise, last_layer_name="1"
         )
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init_precision(laplace, model):
-    feature_extractor = FeatureExtractor(model, last_layer_name='1')
+    feature_extractor = FeatureExtractor(model, last_layer_name="1")
     model_params = list(feature_extractor.last_layer.parameters())
-    setattr(model, 'n_layers', 2)  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", 2)  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     # float
     precision = 10.6
-    lap = laplace(
-        model, likelihood='regression', prior_precision=precision, last_layer_name='1'
+    laplace(
+        model, likelihood="regression", prior_precision=precision, last_layer_name="1"
     )
     # torch.tensor 0-dim
     precision = torch.tensor(10.6)
-    lap = laplace(
-        model, likelihood='regression', prior_precision=precision, last_layer_name='1'
+    laplace(
+        model, likelihood="regression", prior_precision=precision, last_layer_name="1"
     )
     # torch.tensor 1-dim
     precision = torch.tensor(10.7).reshape(-1)
-    lap = laplace(
-        model, likelihood='regression', prior_precision=precision, last_layer_name='1'
+    laplace(
+        model, likelihood="regression", prior_precision=precision, last_layer_name="1"
     )
+
     # torch.tensor 1-dim param-shape
     if laplace == KronLLLaplace:  # kron only supports per layer
         with pytest.raises(ValueError):
             precision = torch.tensor(10.7).reshape(-1).repeat(model.n_params)
-            lap = laplace(
+            laplace(
                 model,
-                likelihood='regression',
+                likelihood="regression",
                 prior_precision=precision,
-                last_layer_name='1',
+                last_layer_name="1",
             )
     else:
         precision = torch.tensor(10.7).reshape(-1).repeat(model.n_params)
-        lap = laplace(
+        laplace(
             model,
-            likelihood='regression',
+            likelihood="regression",
             prior_precision=precision,
-            last_layer_name='1',
+            last_layer_name="1",
         )
     # torch.tensor 1-dim layer-shape
     precision = torch.tensor(10.7).reshape(-1).repeat(model.n_layers)
-    lap = laplace(
-        model, likelihood='regression', prior_precision=precision, last_layer_name='1'
+    laplace(
+        model, likelihood="regression", prior_precision=precision, last_layer_name="1"
     )
 
     # other than that should fail
     # higher dim
     precision = torch.tensor(10.6).reshape(1, 1)
     with pytest.raises(ValueError):
-        lap = laplace(
+        laplace(
             model,
-            likelihood='regression',
+            likelihood="regression",
             prior_precision=precision,
-            last_layer_name='1',
+            last_layer_name="1",
         )
     # unmatched dim
     precision = torch.tensor(10.6).reshape(-1).repeat(17)
     with pytest.raises(ValueError):
         laplace(
             model,
-            likelihood='regression',
+            likelihood="regression",
             prior_precision=precision,
-            last_layer_name='1',
+            last_layer_name="1",
         )
     # other datatype, only reals supported
-    precision = '1.5'
+    precision = "1.5"
     with pytest.raises(ValueError):
         laplace(
             model,
-            likelihood='regression',
+            likelihood="regression",
             prior_precision=precision,
-            last_layer_name='1',
+            last_layer_name="1",
         )
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init_prior_mean_and_scatter(laplace, model, class_loader):
     lap_scalar_mean = laplace(
         model,
-        'classification',
-        last_layer_name='1',
+        "classification",
+        last_layer_name="1",
         prior_precision=1e-2,
         prior_mean=1.0,
     )
     assert torch.allclose(lap_scalar_mean.prior_mean, torch.tensor([1.0]))
     lap_tensor_mean = laplace(
         model,
-        'classification',
-        last_layer_name='1',
+        "classification",
+        last_layer_name="1",
         prior_precision=1e-2,
         prior_mean=torch.ones(1),
     )
     assert torch.allclose(lap_tensor_mean.prior_mean, torch.tensor([1.0]))
     lap_tensor_scalar_mean = laplace(
         model,
-        'classification',
-        last_layer_name='1',
+        "classification",
+        last_layer_name="1",
         prior_precision=1e-2,
         prior_mean=torch.ones(1)[0],
     )
     assert torch.allclose(lap_tensor_scalar_mean.prior_mean, torch.tensor(1.0))
     lap_tensor_full_mean = laplace(
         model,
-        'classification',
-        last_layer_name='1',
+        "classification",
+        last_layer_name="1",
         prior_precision=1e-2,
         prior_mean=torch.ones(20 * 2 + 2),
     )
@@ -306,8 +307,8 @@ def test_laplace_init_prior_mean_and_scatter(laplace, model, class_loader):
         prior_mean = torch.ones(20 * 2 + 2).unsqueeze(-1)
         laplace(
             model,
-            'classification',
-            last_layer_name='1',
+            "classification",
+            last_layer_name="1",
             prior_precision=1e-2,
             prior_mean=prior_mean,
         )
@@ -317,8 +318,8 @@ def test_laplace_init_prior_mean_and_scatter(laplace, model, class_loader):
         prior_mean = torch.ones(20 * 2 - 3)
         laplace(
             model,
-            'classification',
-            last_layer_name='1',
+            "classification",
+            last_layer_name="1",
             prior_precision=1e-2,
             prior_mean=prior_mean,
         )
@@ -327,28 +328,28 @@ def test_laplace_init_prior_mean_and_scatter(laplace, model, class_loader):
     with pytest.raises(ValueError):
         laplace(
             model,
-            'classification',
-            last_layer_name='1',
+            "classification",
+            last_layer_name="1",
             prior_precision=1e-2,
-            prior_mean='72',
+            prior_mean="72",
         )
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_laplace_init_temperature(laplace, model):
     # valid float
     T = 1.1
     lap = laplace(
-        model, likelihood='classification', temperature=T, last_layer_name='1'
+        model, likelihood="classification", temperature=T, last_layer_name="1"
     )
     assert lap.temperature == T
 
 
 @pytest.mark.parametrize(
-    'laplace,lh', product(flavors, ['classification', 'regression'])
+    "laplace,lh", product(flavors, ["classification", "regression"])
 )
-@pytest.mark.parametrize('multidim', [False, True])
-@pytest.mark.parametrize('reduction', [f.value for f in FeatureReduction] + [None])
+@pytest.mark.parametrize("multidim", [False, True])
+@pytest.mark.parametrize("reduction", [f.value for f in FeatureReduction] + [None])
 def test_laplace_functionality(
     laplace,
     lh,
@@ -361,7 +362,7 @@ def test_laplace_functionality(
     class_loader,
     multidim_class_loader,
 ):
-    if lh == 'classification':
+    if lh == "classification":
         loader = class_loader if not multidim else multidim_class_loader
         sigma_noise = 1.0
     else:
@@ -369,11 +370,10 @@ def test_laplace_functionality(
         sigma_noise = 0.3
 
     if not multidim:
-        model = model
-        last_layer_name = '1'
+        last_layer_name = "1"
     else:
         model = model_with_reduction
-        last_layer_name = 'fc2'
+        last_layer_name = "fc2"
 
     lap = laplace(
         model,
@@ -394,7 +394,7 @@ def test_laplace_functionality(
     # Test log likelihood (Train)
     log_lik = lap.log_likelihood
     # compute true log lik
-    if lh == 'classification':
+    if lh == "classification":
         log_lik_true = Categorical(logits=f).log_prob(y).sum()
         assert torch.allclose(log_lik, log_lik_true)
     else:
@@ -446,7 +446,7 @@ def test_laplace_functionality(
         Sigma = torch.diag(lap.posterior_variance)
     _, phi = feature_extractor.forward_with_features(X)
     Js, f = jacobians_naive(feature_extractor.last_layer, phi)
-    true_f_var = torch.einsum('mkp,pq,mcq->mkc', Js, Sigma, Js)
+    true_f_var = torch.einsum("mkp,pq,mcq->mkc", Js, Sigma, Js)
     # test last-layer Jacobians
     comp_Js, comp_f = lap.backend.last_layer_jacobians(X)
     assert torch.allclose(Js, comp_Js)
@@ -455,121 +455,121 @@ def test_laplace_functionality(
     assert torch.allclose(true_f_var, comp_f_var, rtol=1e-4)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_regression_predictive(laplace, model, reg_loader):
-    lap = laplace(model, 'regression', sigma_noise=0.3, prior_precision=0.7)
+    lap = laplace(model, "regression", sigma_noise=0.3, prior_precision=0.7)
     lap.fit(reg_loader)
     X, y = reg_loader.dataset.tensors
     f = model(X)
 
     # error
     with pytest.raises(ValueError):
-        lap(X, pred_type='linear')
+        lap(X, pred_type="linear")
 
     # GLM predictive, functional variance tested already above.
-    f_mu, f_var = lap(X, pred_type='glm')
+    f_mu, f_var = lap(X, pred_type="glm")
     assert torch.allclose(f_mu, f)
     assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1], f_mu.shape[1]])
     assert len(f_mu) == len(X)
 
     # NN predictive (only diagonal variance estimation)
-    f_mu, f_var = lap(X, pred_type='nn', link_approx='mc')
+    f_mu, f_var = lap(X, pred_type="nn", link_approx="mc")
     assert f_mu.shape == f_var.shape
     assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1]])
     assert len(f_mu) == len(X)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_classification_predictive(laplace, model, class_loader):
-    lap = laplace(model, 'classification', prior_precision=0.7)
+    lap = laplace(model, "classification", prior_precision=0.7)
     lap.fit(class_loader)
     X, y = class_loader.dataset.tensors
     f = torch.softmax(model(X), dim=-1)
 
     # error
     with pytest.raises(ValueError):
-        lap(X, pred_type='linear')
+        lap(X, pred_type="linear")
 
     # GLM predictive
-    f_pred = lap(X, pred_type='glm', link_approx='mc', n_samples=100)
+    f_pred = lap(X, pred_type="glm", link_approx="mc", n_samples=100)
     assert f_pred.shape == f.shape
     assert torch.allclose(
         f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
     )  # sum up to 1
-    f_pred = lap(X, pred_type='glm', link_approx='probit')
+    f_pred = lap(X, pred_type="glm", link_approx="probit")
     assert f_pred.shape == f.shape
     assert torch.allclose(
         f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
     )  # sum up to 1
-    f_pred = lap(X, pred_type='glm', link_approx='bridge')
+    f_pred = lap(X, pred_type="glm", link_approx="bridge")
     assert f_pred.shape == f.shape
     assert torch.allclose(
         f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
     )  # sum up to 1
-    f_pred = lap(X, pred_type='glm', link_approx='bridge_norm')
+    f_pred = lap(X, pred_type="glm", link_approx="bridge_norm")
     assert f_pred.shape == f.shape
     assert torch.allclose(
         f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
     )  # sum up to 1
 
     # NN predictive
-    f_pred = lap(X, pred_type='nn', link_approx='mc', n_samples=100)
+    f_pred = lap(X, pred_type="nn", link_approx="mc", n_samples=100)
     assert f_pred.shape == f.shape
     assert torch.allclose(
         f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
     )  # sum up to 1
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_regression_predictive_samples(laplace, model, reg_loader):
-    lap = laplace(model, 'regression', sigma_noise=0.3, prior_precision=0.7)
+    lap = laplace(model, "regression", sigma_noise=0.3, prior_precision=0.7)
     lap.fit(reg_loader)
     X, y = reg_loader.dataset.tensors
     f = model(X)
 
     # error
     with pytest.raises(ValueError):
-        lap(X, pred_type='linear')
+        lap(X, pred_type="linear")
 
     # GLM predictive, functional variance tested already above.
-    fsamples = lap.predictive_samples(X, pred_type='glm', n_samples=100)
+    fsamples = lap.predictive_samples(X, pred_type="glm", n_samples=100)
     assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
 
     # NN predictive (only diagonal variance estimation)
-    fsamples = lap.predictive_samples(X, pred_type='nn', n_samples=100)
+    fsamples = lap.predictive_samples(X, pred_type="nn", n_samples=100)
     assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_classification_predictive_samples(laplace, model, class_loader):
-    lap = laplace(model, 'classification', prior_precision=0.7)
+    lap = laplace(model, "classification", prior_precision=0.7)
     lap.fit(class_loader)
     X, y = class_loader.dataset.tensors
     f = torch.softmax(model(X), dim=-1)
 
     # error
     with pytest.raises(ValueError):
-        lap(X, pred_type='linear')
+        lap(X, pred_type="linear")
 
     # GLM predictive
-    fsamples = lap.predictive_samples(X, pred_type='glm', n_samples=100)
+    fsamples = lap.predictive_samples(X, pred_type="glm", n_samples=100)
     assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
     assert np.allclose(fsamples.sum().item(), len(f) * 100)  # sum up to 1
 
     # NN predictive
-    fsamples = lap.predictive_samples(X, pred_type='nn', n_samples=100)
+    fsamples = lap.predictive_samples(X, pred_type="nn", n_samples=100)
     assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
     assert np.allclose(fsamples.sum().item(), len(f) * 100)  # sum up to 1
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_backprop_glm(laplace, model, reg_loader):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True)
+    lap = laplace(model, "regression", enable_backprop=True)
     lap.fit(reg_loader)
-    f_mu, f_var = lap(X, pred_type='glm')
+    f_mu, f_var = lap(X, pred_type="glm")
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
@@ -581,14 +581,14 @@ def test_backprop_glm(laplace, model, reg_loader):
         assert False
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_backprop_glm_joint(laplace, model, reg_loader):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True)
+    lap = laplace(model, "regression", enable_backprop=True)
     lap.fit(reg_loader)
-    f_mu, f_cov = lap(X, pred_type='glm', joint=True)
+    f_mu, f_cov = lap(X, pred_type="glm", joint=True)
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
@@ -600,14 +600,14 @@ def test_backprop_glm_joint(laplace, model, reg_loader):
         assert False
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_backprop_glm_mc(laplace, model, reg_loader):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True)
+    lap = laplace(model, "regression", enable_backprop=True)
     lap.fit(reg_loader)
-    f_mu, f_var = lap(X, pred_type='glm', link_approx='mc')
+    f_mu, f_var = lap(X, pred_type="glm", link_approx="mc")
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
@@ -619,14 +619,14 @@ def test_backprop_glm_mc(laplace, model, reg_loader):
         assert False
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_backprop_nn(laplace, model, reg_loader):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True)
+    lap = laplace(model, "regression", enable_backprop=True)
     lap.fit(reg_loader)
-    f_mu, f_var = lap(X, pred_type='nn', link_approx='mc', n_samples=10)
+    f_mu, f_var = lap(X, pred_type="nn", link_approx="mc", n_samples=10)
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index 757130c..e1701ae 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -1,14 +1,13 @@
-import pytest
 import numpy as np
+import pytest
 import torch
 from torch import nn
 from torch.nn.utils import parameters_to_vector
 
+from laplace.curvature import AsdlGGN, BackPackGGN
 from laplace.utils import Kron, block_diag
 from laplace.utils import kron as kron_prod
-from laplace.curvature import BackPackGGN, AsdlGGN
-from tests.utils import get_psd_matrix, get_diag_psd_matrix, jacobians_naive
-
+from tests.utils import get_diag_psd_matrix, get_psd_matrix, jacobians_naive
 
 torch.set_default_tensor_type(torch.DoubleTensor)
 
@@ -16,23 +15,23 @@
 @pytest.fixture
 def model():
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
 @pytest.fixture
 def small_model():
     model = torch.nn.Sequential(nn.Linear(3, 5), nn.Tanh(), nn.Linear(5, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     return model
 
 
 def test_init_from_model(model):
-    kron = Kron.init_from_model(model, 'cpu')
-    expected_sizes = [[20*20, 3*3], [20*20], [2*2, 20*20], [2*2]]
+    kron = Kron.init_from_model(model, "cpu")
+    expected_sizes = [[20 * 20, 3 * 3], [20 * 20], [2 * 2, 20 * 20], [2 * 2]]
     for facs, exp_facs in zip(kron.kfacs, expected_sizes):
         for fi, exp_fi in zip(facs, exp_facs):
             assert torch.all(fi == 0)
@@ -40,8 +39,8 @@ def test_init_from_model(model):
 
 
 def test_init_from_iterable(model):
-    kron = Kron.init_from_model(model.parameters(), 'cpu')
-    expected_sizes = [[20*20, 3*3], [20*20], [2*2, 20*20], [2*2]]
+    kron = Kron.init_from_model(model.parameters(), "cpu")
+    expected_sizes = [[20 * 20, 3 * 3], [20 * 20], [2 * 2, 20 * 20], [2 * 2]]
     for facs, exp_facs in zip(kron.kfacs, expected_sizes):
         for fi, exp_fi in zip(facs, exp_facs):
             assert torch.all(fi == 0)
@@ -49,7 +48,7 @@ def test_init_from_iterable(model):
 
 
 def test_addition(model):
-    kron = Kron.init_from_model(model.parameters(), 'cpu')
+    kron = Kron.init_from_model(model.parameters(), "cpu")
     expected_sizes = [[20, 3], [20], [2, 20], [2]]
     to_add = Kron([[torch.randn(i, i) for i in sizes] for sizes in expected_sizes])
     kron += to_add
@@ -81,10 +80,12 @@ def test_decompose():
     kfacs = [[get_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
     kron = Kron(kfacs)
     kron_decomp = kron.decompose()
-    for facs, Qs, ls in zip(kron.kfacs, kron_decomp.eigenvectors, kron_decomp.eigenvalues):
+    for facs, Qs, ls in zip(
+        kron.kfacs, kron_decomp.eigenvectors, kron_decomp.eigenvalues
+    ):
         if len(facs) == 1:
-            H, Q, l = facs[0], Qs[0], ls[0]
-            reconstructed = Q @ torch.diag(l) @ Q.T
+            H, Q, eigval = facs[0], Qs[0], ls[0]
+            reconstructed = Q @ torch.diag(eigval) @ Q.T
             assert torch.allclose(H, reconstructed, rtol=1e-3)
         if len(facs) == 2:
             gtruth = kron_prod(facs[0], facs[1])
@@ -100,10 +101,12 @@ def test_decompose():
     diag_kfacs = [[get_diag_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
     kron = Kron(diag_kfacs)
     kron_decomp = kron.decompose()
-    for facs, Qs, ls in zip(kron.kfacs, kron_decomp.eigenvectors, kron_decomp.eigenvalues):
+    for facs, Qs, ls in zip(
+        kron.kfacs, kron_decomp.eigenvectors, kron_decomp.eigenvalues
+    ):
         if len(facs) == 1:
-            H, Q, l = facs[0], Qs[0], ls[0]
-            reconstructed = (Q @ torch.diag(l) @ Q.T).diag()
+            H, Q, eigval = facs[0], Qs[0], ls[0]
+            reconstructed = (Q @ torch.diag(eigval) @ Q.T).diag()
             assert torch.allclose(H, reconstructed, rtol=1e-3)
         if len(facs) == 2:
             gtruth = kron_prod(facs[0].diag(), facs[1].diag())
@@ -137,7 +140,7 @@ def test_bmm_dense(small_model):
     y = torch.randn(5, 2)
 
     # Dense Kronecker factors.
-    backend = BackPackGGN(model, 'regression', stochastic=False)
+    backend = BackPackGGN(model, "regression", stochastic=False)
     _, kron = backend.kron(X, y, N=5)
     kron_decomp = kron.decompose()
     Js, _ = jacobians_naive(model, X)
@@ -154,29 +157,29 @@ def test_bmm_dense(small_model):
     # test J @ Kron_decomp @ Jt (square form)
     JS = kron_decomp.bmm(Js, exponent=1)
     JS_true = Js @ S
-    JSJ_true = torch.bmm(JS_true, Js.transpose(1,2))
-    JSJ = torch.bmm(JS, Js.transpose(1,2))
+    JSJ_true = torch.bmm(JS_true, Js.transpose(1, 2))
+    JSJ = torch.bmm(JS, Js.transpose(1, 2))
     assert torch.allclose(JSJ, JSJ_true)
     assert torch.allclose(JS, JS_true)
 
     # test J @ Kron @ Jt (square form)
     JS_nodecomp = kron.bmm(Js)
-    JSJ_nodecomp = torch.bmm(JS_nodecomp, Js.transpose(1,2))
+    JSJ_nodecomp = torch.bmm(JS_nodecomp, Js.transpose(1, 2))
     assert torch.allclose(JSJ_nodecomp, JSJ)
     assert torch.allclose(JS_nodecomp, JS)
 
     # test J @ S_inv @ J (funcitonal variance)
     JSJ = kron_decomp.inv_square_form(Js)
     S_inv = S.inverse()
-    JSJ_true = torch.bmm(Js @ S_inv, Js.transpose(1,2))
+    JSJ_true = torch.bmm(Js @ S_inv, Js.transpose(1, 2))
     assert torch.allclose(JSJ, JSJ_true)
 
     # test J @ S^-1/2  (sampling)
-    JS = kron_decomp.bmm(Js, exponent=-1/2)
-    JSJ = torch.bmm(JS, Js.transpose(1,2))
-    l, Q = torch.linalg.eigh(S_inv, UPLO='U')
-    JS_true = Js @ Q @ torch.diag(torch.sqrt(l)) @ Q.T
-    JSJ_true = torch.bmm(JS_true, Js.transpose(1,2))
+    JS = kron_decomp.bmm(Js, exponent=-1 / 2)
+    JSJ = torch.bmm(JS, Js.transpose(1, 2))
+    eigval, Q = torch.linalg.eigh(S_inv, UPLO="U")
+    JS_true = Js @ Q @ torch.diag(torch.sqrt(eigval)) @ Q.T
+    JSJ_true = torch.bmm(JS_true, Js.transpose(1, 2))
     assert torch.allclose(JS, JS_true)
     assert torch.allclose(JSJ, JSJ_true)
 
@@ -195,7 +198,7 @@ def test_bmm_dense(small_model):
     assert torch.allclose(JS, JS_nodecomp)
 
 
-@pytest.mark.skip(reason='For compatibility with ASDL master')
+@pytest.mark.skip(reason="For compatibility with ASDL master")
 def test_bmm_diag(small_model):
     model = small_model
     # model = single_output_model
@@ -203,7 +206,7 @@ def test_bmm_diag(small_model):
     y = torch.randn(5, 2)
 
     # Diagonal Kronecker factors.
-    backend = AsdlGGN(model, 'regression', stochastic=False)
+    backend = AsdlGGN(model, "regression", stochastic=False)
     _, kron = backend.kron(X, y, N=5, diag_A=True, diag_B=True)
     kron_decomp = kron.decompose()
     Js, _ = jacobians_naive(model, X)
@@ -222,29 +225,29 @@ def test_bmm_diag(small_model):
     # test J @ Kron_decomp @ Jt (square form)
     JS = kron_decomp.bmm(Js, exponent=1)
     JS_true = Js @ S
-    JSJ_true = torch.bmm(JS_true, Js.transpose(1,2))
-    JSJ = torch.bmm(JS, Js.transpose(1,2))
+    JSJ_true = torch.bmm(JS_true, Js.transpose(1, 2))
+    JSJ = torch.bmm(JS, Js.transpose(1, 2))
     assert torch.allclose(JSJ, JSJ_true)
     assert torch.allclose(JS, JS_true)
 
     # test J @ Kron @ Jt (square form)
     JS_nodecomp = kron.bmm(Js)
-    JSJ_nodecomp = torch.bmm(JS_nodecomp, Js.transpose(1,2))
+    JSJ_nodecomp = torch.bmm(JS_nodecomp, Js.transpose(1, 2))
     assert torch.allclose(JSJ_nodecomp, JSJ)
     assert torch.allclose(JS_nodecomp, JS)
 
     # test J @ S_inv @ J (funcitonal variance)
     JSJ = kron_decomp.inv_square_form(Js)
     S_inv = S.inverse()
-    JSJ_true = torch.bmm(Js @ S_inv, Js.transpose(1,2))
+    JSJ_true = torch.bmm(Js @ S_inv, Js.transpose(1, 2))
     assert torch.allclose(JSJ, JSJ_true)
 
     # test J @ S^-1/2  (sampling)
-    JS = kron_decomp.bmm(Js, exponent=-1/2)
-    JSJ = torch.bmm(JS, Js.transpose(1,2))
-    l, Q = torch.linalg.eigh(S_inv, UPLO='U')
-    JS_true = Js @ Q @ torch.diag(torch.sqrt(l)) @ Q.T
-    JSJ_true = torch.bmm(JS_true, Js.transpose(1,2))
+    JS = kron_decomp.bmm(Js, exponent=-1 / 2)
+    JSJ = torch.bmm(JS, Js.transpose(1, 2))
+    eigval, Q = torch.linalg.eigh(S_inv, UPLO="U")
+    JS_true = Js @ Q @ torch.diag(torch.sqrt(eigval)) @ Q.T
+    JSJ_true = torch.bmm(JS_true, Js.transpose(1, 2))
     assert torch.allclose(JS, JS_true)
     assert torch.allclose(JSJ, JSJ_true)
 
@@ -271,7 +274,9 @@ def test_matrix_consistent():
     kron = Kron(kfacs)
     kron_decomp = kron.decompose()
     assert torch.allclose(kron.to_matrix(), kron_decomp.to_matrix(exponent=1))
-    assert torch.allclose(kron.to_matrix().inverse(), kron_decomp.to_matrix(exponent=-1))
+    assert torch.allclose(
+        kron.to_matrix().inverse(), kron_decomp.to_matrix(exponent=-1)
+    )
     M_true = kron.to_matrix()
     M_true.diagonal().add_(3.4)
     kron_decomp += torch.tensor(3.4)
@@ -281,7 +286,9 @@ def test_matrix_consistent():
     kron = Kron(diag_kfacs)
     kron_decomp = kron.decompose()
     assert torch.allclose(kron.to_matrix(), kron_decomp.to_matrix(exponent=1))
-    assert torch.allclose(kron.to_matrix().inverse(), kron_decomp.to_matrix(exponent=-1))
+    assert torch.allclose(
+        kron.to_matrix().inverse(), kron_decomp.to_matrix(exponent=-1)
+    )
     M_true = kron.to_matrix()
     M_true.diagonal().add_(3.4)
     kron_decomp += torch.tensor(3.4)
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index 2edb72e..02d9683 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -1,7 +1,9 @@
+import math
+
 import torch
 from torch.nn import functional as F
+
 from laplace.utils import RunningNLLMetric
-import math
 
 
 def test_running_nll_metric():
@@ -18,7 +20,9 @@ def test_running_nll_metric():
     all_probs, all_targets = torch.cat(all_probs, 0), torch.cat(all_targets, 0)
 
     nll_running = metric.compute().item()
-    nll_offline = F.nll_loss(all_probs.log().flatten(end_dim=-2), all_targets.flatten()).item()
+    nll_offline = F.nll_loss(
+        all_probs.log().flatten(end_dim=-2), all_targets.flatten()
+    ).item()
 
     assert math.isclose(nll_running, nll_offline, rel_tol=1e-7)
 
@@ -40,7 +44,11 @@ def test_running_nll_metric_ignore_idx():
     all_probs, all_targets = torch.cat(all_probs, 0), torch.cat(all_targets, 0)
 
     nll_running = metric.compute().item()
-    nll_offline = F.nll_loss(all_probs.log().flatten(end_dim=-2), all_targets.flatten(), ignore_index=ignore_idx).item()
+    nll_offline = F.nll_loss(
+        all_probs.log().flatten(end_dim=-2),
+        all_targets.flatten(),
+        ignore_index=ignore_idx,
+    ).item()
 
     print(nll_running, nll_offline)
 
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
index 612a433..f34c097 100644
--- a/tests/test_serialization.py
+++ b/tests/test_serialization.py
@@ -1,33 +1,42 @@
-import pytest
+import os
+from collections import OrderedDict
 from importlib.util import find_spec
+
+import pytest
 import torch
 from torch import nn
 from torch.nn.utils import parameters_to_vector
 from torch.utils.data import DataLoader, TensorDataset
-import os
-from collections import OrderedDict
 
-from laplace import Laplace
-from laplace.laplace import (
-    FullLaplace,
-    KronLaplace,
+from laplace import (
     DiagLaplace,
-    LowRankLaplace,
-    FullLLLaplace,
-    KronLLLaplace,
     DiagLLLaplace,
-    LLLaplace,
     DiagSubnetLaplace,
+    FullLaplace,
+    FullLLLaplace,
     FullSubnetLaplace,
+    KronLaplace,
+    KronLLLaplace,
+    Laplace,
+    LLLaplace,
+    LowRankLaplace,
     SubnetLaplace,
 )
 
-
-
 torch.manual_seed(240)
 torch.set_default_tensor_type(torch.DoubleTensor)
-lrlaplace_param = pytest.param(LowRankLaplace, marks=pytest.mark.xfail(reason='Unimplemented in the new ASDL'))
-flavors = [FullLaplace, KronLaplace, DiagLaplace, lrlaplace_param, FullLLLaplace, KronLLLaplace, DiagLLLaplace]
+lrlaplace_param = pytest.param(
+    LowRankLaplace, marks=pytest.mark.xfail(reason="Unimplemented in the new ASDL")
+)
+flavors = [
+    FullLaplace,
+    KronLaplace,
+    DiagLaplace,
+    lrlaplace_param,
+    FullLLLaplace,
+    KronLLLaplace,
+    DiagLLLaplace,
+]
 flavors_no_llla = [FullLaplace, KronLaplace, DiagLaplace, lrlaplace_param]
 flavors_llla = [FullLLLaplace, KronLLLaplace, DiagLLLaplace]
 flavors_subnet = [DiagSubnetLaplace, FullSubnetLaplace]
@@ -37,35 +46,32 @@
 @pytest.fixture
 def model():
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
 @pytest.fixture
 def model2():
     model = torch.nn.Sequential(nn.Linear(3, 25), nn.Linear(25, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
 @pytest.fixture
 def model3():
     model = torch.nn.Sequential(
-        OrderedDict([
-            ('fc1', nn.Linear(3, 20)),
-            ('clf', nn.Linear(20, 2))
-        ])
+        OrderedDict([("fc1", nn.Linear(3, 20)), ("clf", nn.Linear(20, 2))])
     )
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -80,20 +86,20 @@ def reg_loader():
 def cleanup():
     yield
     # Run after test
-    if os.path.exists('state_dict.bin'):
-        os.remove('state_dict.bin')
+    if os.path.exists("state_dict.bin"):
+        os.remove("state_dict.bin")
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_serialize(laplace, model, reg_loader):
-    la = laplace(model, 'regression')
+    la = laplace(model, "regression")
     la.fit(reg_loader)
     la.optimize_prior_precision()
     la.sigma_noise = 1231
-    torch.save(la.state_dict(), 'state_dict.bin')
+    torch.save(la.state_dict(), "state_dict.bin")
 
-    la2 = laplace(model, 'regression')
-    la2.load_state_dict(torch.load('state_dict.bin'))
+    la2 = laplace(model, "regression")
+    la2.load_state_dict(torch.load("state_dict.bin"))
 
     assert la.sigma_noise == la2.sigma_noise
 
@@ -104,17 +110,17 @@ def test_serialize(laplace, model, reg_loader):
     assert torch.allclose(f_var, f_var2)
 
 
-@pytest.mark.parametrize('laplace', flavors_no_llla[:-1])
+@pytest.mark.parametrize("laplace", flavors_no_llla[:-1])
 def test_serialize_override(laplace, model, reg_loader):
-    la = laplace(model, 'regression')
+    la = laplace(model, "regression")
     la.fit(reg_loader)
     la.optimize_prior_precision()
     la.sigma_noise = 1231
     H_orig = la.H_facs.to_matrix() if laplace == KronLaplace else la.H
-    torch.save(la.state_dict(), 'state_dict.bin')
+    torch.save(la.state_dict(), "state_dict.bin")
 
-    la2 = laplace(model, 'regression')
-    la2.load_state_dict(torch.load('state_dict.bin'))
+    la2 = laplace(model, "regression")
+    la2.load_state_dict(torch.load("state_dict.bin"))
 
     # Emulating continual learning
     la2.fit(reg_loader, override=False)
@@ -123,14 +129,14 @@ def test_serialize_override(laplace, model, reg_loader):
     assert torch.allclose(2 * H_orig, H_new)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_serialize_no_pickle(laplace, model, reg_loader):
-    la = laplace(model, 'regression')
+    la = laplace(model, "regression")
     la.fit(reg_loader)
     la.optimize_prior_precision()
     la.sigma_noise = 1231
-    torch.save(la.state_dict(), 'state_dict.bin')
-    state_dict = torch.load('state_dict.bin')
+    torch.save(la.state_dict(), "state_dict.bin")
+    state_dict = torch.load("state_dict.bin")
 
     # Make sure no pickle object
     for val in state_dict.values():
@@ -138,17 +144,17 @@ def test_serialize_no_pickle(laplace, model, reg_loader):
             assert isinstance(val, (list, tuple, int, float, str, bool, torch.Tensor))
 
 
-@pytest.mark.parametrize('laplace', flavors_subnet)
+@pytest.mark.parametrize("laplace", flavors_subnet)
 def test_serialize_subnetlaplace(laplace, model, reg_loader):
     subnetwork_indices = torch.LongTensor([1, 10, 104, 44])
-    la = laplace(model, 'regression', subnetwork_indices=subnetwork_indices)
+    la = laplace(model, "regression", subnetwork_indices=subnetwork_indices)
     la.fit(reg_loader)
     la.optimize_prior_precision()
     la.sigma_noise = 1231
-    torch.save(la.state_dict(), 'state_dict.bin')
+    torch.save(la.state_dict(), "state_dict.bin")
 
-    la2 = laplace(model, 'regression', subnetwork_indices=subnetwork_indices)
-    la2.load_state_dict(torch.load('state_dict.bin'))
+    la2 = laplace(model, "regression", subnetwork_indices=subnetwork_indices)
+    la2.load_state_dict(torch.load("state_dict.bin"))
 
     assert la.sigma_noise == la2.sigma_noise
 
@@ -159,107 +165,113 @@ def test_serialize_subnetlaplace(laplace, model, reg_loader):
     assert torch.allclose(f_var, f_var2)
 
 
-@pytest.mark.parametrize('laplace', flavors_no_llla)
+@pytest.mark.parametrize("laplace", flavors_no_llla)
 def test_serialize_fail_different_models(laplace, model, model2, reg_loader):
-    la = laplace(model, 'regression')
+    la = laplace(model, "regression")
     la.fit(reg_loader)
     la.optimize_prior_precision()
     la.sigma_noise = 1231
-    torch.save(la.state_dict(), 'state_dict.bin')
+    torch.save(la.state_dict(), "state_dict.bin")
 
-    la2 = laplace(model2, 'regression')
+    la2 = laplace(model2, "regression")
 
     with pytest.raises(ValueError):
-        la2.load_state_dict(torch.load('state_dict.bin'))
+        la2.load_state_dict(torch.load("state_dict.bin"))
 
 
 def test_serialize_fail_different_hess_structures(model, reg_loader):
-    la = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='kron')
+    la = Laplace(model, "regression", subset_of_weights="all", hessian_structure="kron")
     la.fit(reg_loader)
     la.optimize_prior_precision()
     la.sigma_noise = 1231
-    torch.save(la.state_dict(), 'state_dict.bin')
+    torch.save(la.state_dict(), "state_dict.bin")
 
-    la2 = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='diag')
+    la2 = Laplace(
+        model, "regression", subset_of_weights="all", hessian_structure="diag"
+    )
 
     with pytest.raises(ValueError):
-        la2.load_state_dict(torch.load('state_dict.bin'))
+        la2.load_state_dict(torch.load("state_dict.bin"))
 
 
 def test_serialize_fail_different_subset_of_weights(model, reg_loader):
-    la = Laplace(model, 'regression', subset_of_weights='last_layer', hessian_structure='diag')
+    la = Laplace(
+        model, "regression", subset_of_weights="last_layer", hessian_structure="diag"
+    )
     la.fit(reg_loader)
     la.optimize_prior_precision()
     la.sigma_noise = 1231
-    torch.save(la.state_dict(), 'state_dict.bin')
+    torch.save(la.state_dict(), "state_dict.bin")
 
-    la2 = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='diag')
+    la2 = Laplace(
+        model, "regression", subset_of_weights="all", hessian_structure="diag"
+    )
 
     with pytest.raises(ValueError):
-        la2.load_state_dict(torch.load('state_dict.bin'))
+        la2.load_state_dict(torch.load("state_dict.bin"))
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_serialize_fail_different_liks(laplace, model, reg_loader):
-    la = laplace(model, 'regression')
+    la = laplace(model, "regression")
     la.fit(reg_loader)
     la.optimize_prior_precision()
     la.sigma_noise = 1231
-    torch.save(la.state_dict(), 'state_dict.bin')
+    torch.save(la.state_dict(), "state_dict.bin")
 
-    la2 = laplace(model, 'classification')
+    la2 = laplace(model, "classification")
 
     with pytest.raises(ValueError):
-        la2.load_state_dict(torch.load('state_dict.bin'))
+        la2.load_state_dict(torch.load("state_dict.bin"))
 
 
-@pytest.mark.parametrize('laplace', flavors_llla)
-def test_serialize_fail_llla_different_last_layer_name(laplace, model, model3, reg_loader):
+@pytest.mark.parametrize("laplace", flavors_llla)
+def test_serialize_fail_llla_different_last_layer_name(
+    laplace, model, model3, reg_loader
+):
     print([n for n, _ in model.named_parameters()])
-    la = laplace(model, 'regression', last_layer_name='1')
+    la = laplace(model, "regression", last_layer_name="1")
     la.fit(reg_loader)
     la.optimize_prior_precision()
     la.sigma_noise = 1231
-    torch.save(la.state_dict(), 'state_dict.bin')
+    torch.save(la.state_dict(), "state_dict.bin")
 
-    la2 = laplace(model3, 'classification', last_layer_name='clf')
+    la2 = laplace(model3, "classification", last_layer_name="clf")
 
     with pytest.raises(ValueError):
-        la2.load_state_dict(torch.load('state_dict.bin'))
+        la2.load_state_dict(torch.load("state_dict.bin"))
 
 
 @pytest.mark.parametrize(
-    'model_device_str',
-    ['cpu', 'cuda'] if torch.cuda.is_available() else ['cpu'],
+    "model_device_str",
+    ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"],
 )
 @pytest.mark.parametrize(
-    'map_location',
-    ['cuda', 'cpu'] if torch.cuda.is_available() else ['cpu'],
+    "map_location",
+    ["cuda", "cpu"] if torch.cuda.is_available() else ["cpu"],
 )
-@pytest.mark.parametrize('laplace', flavors + flavors_subnet)
+@pytest.mark.parametrize("laplace", flavors + flavors_subnet)
 def test_map_location(
     laplace, model, reg_loader, tmp_path, map_location, model_device_str
 ):
     # Only for LLLaplace, we need dill.
     #   AttributeError: Can't pickle local object 'FeatureExtractor._get_hook.<locals>.hook'
     if issubclass(laplace, LLLaplace):
-        if find_spec('dill') is None:
-            pytest.skip(
-                reason='dill package not found but needed for this test'
-            )
+        if find_spec("dill") is None:
+            pytest.skip(reason="dill package not found but needed for this test")
         else:
             import dill
 
-            torch_save = lambda obj, fn: torch.save(
-                obj, fn, pickle_module=dill
-            )
+            def torch_save(obj, fn):
+                return torch.save(obj, fn, pickle_module=dill)
     else:
         # Use default pickle_module=pickle, but no need to import pickle here
         # just to set pickle_module=pickle.
-        torch_save = lambda obj, fn: torch.save(obj, fn)
+        def torch_save(obj, fn):
+            return torch.save(obj, fn)
 
     device = torch.device(model_device_str)
-    kwds = dict(model=model.to(device), likelihood='regression')
+    kwds = dict(model=model.to(device), likelihood="regression")
     if issubclass(laplace, SubnetLaplace):
         kwds.update(subnetwork_indices=torch.LongTensor([1, 10, 104, 44]))
     la = laplace(**kwds)
@@ -272,7 +284,7 @@ def test_map_location(
     la.optimize_prior_precision()
     la.sigma_noise = 1231
 
-    save_fn = tmp_path / 'la.pt'
+    save_fn = tmp_path / "la.pt"
     torch_save(la, save_fn)
     la2 = torch.load(save_fn, map_location=map_location)
 
@@ -282,17 +294,13 @@ def test_map_location(
     # This is redundant if *all* will be used in the forward call below, where
     # things would fail if self._device is wrong.
     for name in [
-        'prior_precision_diag',
-        'prior_mean',
-        'prior_precision',
-        'sigma_noise',
+        "prior_precision_diag",
+        "prior_mean",
+        "prior_precision",
+        "sigma_noise",
     ]:
-        assert (
-            getattr(la, name).device.type == device.type
-        ), f'la.{name} failed'
-        assert (
-            getattr(la2, name).device.type == map_location
-        ), f'la2.{name} failed'
+        assert getattr(la, name).device.type == device.type, f"la.{name} failed"
+        assert getattr(la2, name).device.type == map_location, f"la2.{name} failed"
 
     # Test tensor attrs.
     for name, obj in vars(la).items():
@@ -302,17 +310,15 @@ def test_map_location(
         # as in BaseLaplace.fit()). self.X won't be affected by self._device.
         # Since the predict() below passes, we ignore the wrong device here
         # for now.
-        if issubclass(laplace, LLLaplace) and name == 'X':
+        if issubclass(laplace, LLLaplace) and name == "X":
             continue
         if isinstance(obj, torch.Tensor):
-            assert obj.device.type == device.type, f'la.{name} failed'
-            assert (
-                getattr(la2, name).device.type == map_location
-            ), f'la2.{name} failed'
+            assert obj.device.type == device.type, f"la.{name} failed"
+            assert getattr(la2, name).device.type == map_location, f"la2.{name} failed"
 
     assert la.sigma_noise == la2.sigma_noise
     X, _ = next(iter(reg_loader))
     f_mean, f_var = la(X.to(device))
     f_mean2, f_var2 = la2(X.to(map_location))
-    assert torch.allclose(f_mean.to('cpu'), f_mean2.to('cpu'))
-    assert torch.allclose(f_var.to('cpu'), f_var2.to('cpu'))
+    assert torch.allclose(f_mean.to("cpu"), f_mean2.to("cpu"))
+    assert torch.allclose(f_var.to("cpu"), f_var2.to("cpu"))
diff --git a/tests/test_subnetlaplace.py b/tests/test_subnetlaplace.py
index 92c888b..90ab55d 100644
--- a/tests/test_subnetlaplace.py
+++ b/tests/test_subnetlaplace.py
@@ -1,34 +1,44 @@
-import pytest
 from itertools import product
 
+import pytest
 import torch
 from torch import nn
 from torch.nn.utils import parameters_to_vector
 from torch.utils.data import DataLoader, TensorDataset
 from torchvision.models import wide_resnet50_2
 
-from laplace import Laplace, SubnetLaplace, FullSubnetLaplace, DiagSubnetLaplace
+from laplace import DiagSubnetLaplace, FullSubnetLaplace, Laplace, SubnetLaplace
 from laplace.baselaplace import DiagLaplace
-from laplace.utils import (SubnetMask, RandomSubnetMask, LargestMagnitudeSubnetMask,
-                           LargestVarianceDiagLaplaceSubnetMask, LargestVarianceSWAGSubnetMask,
-                           ParamNameSubnetMask, ModuleNameSubnetMask, LastLayerSubnetMask)
-
+from laplace.utils import (
+    LargestMagnitudeSubnetMask,
+    LargestVarianceDiagLaplaceSubnetMask,
+    LargestVarianceSWAGSubnetMask,
+    LastLayerSubnetMask,
+    ModuleNameSubnetMask,
+    ParamNameSubnetMask,
+    RandomSubnetMask,
+    SubnetMask,
+)
 
 torch.manual_seed(240)
 torch.set_default_tensor_type(torch.DoubleTensor)
-score_based_subnet_masks = [RandomSubnetMask, LargestMagnitudeSubnetMask,
-                            LargestVarianceDiagLaplaceSubnetMask, LargestVarianceSWAGSubnetMask]
+score_based_subnet_masks = [
+    RandomSubnetMask,
+    LargestMagnitudeSubnetMask,
+    LargestVarianceDiagLaplaceSubnetMask,
+    LargestVarianceSWAGSubnetMask,
+]
 layer_subnet_masks = [ParamNameSubnetMask, ModuleNameSubnetMask, LastLayerSubnetMask]
 all_subnet_masks = score_based_subnet_masks + layer_subnet_masks
-likelihoods = ['classification', 'regression']
-hessian_structures = ['full', 'diag']
+likelihoods = ["classification", "regression"]
+hessian_structures = ["full", "diag"]
 
 
 @pytest.fixture
 def model():
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
     model_params = list(model.parameters())
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
     return model
 
 
@@ -52,7 +62,7 @@ def reg_loader():
     return DataLoader(TensorDataset(X, y), batch_size=3)
 
 
-@pytest.mark.parametrize('likelihood', likelihoods)
+@pytest.mark.parametrize("likelihood", likelihoods)
 def test_subnet_laplace_init(model, likelihood):
     # use random subnet mask for this test
     subnetwork_mask = RandomSubnetMask
@@ -61,34 +71,60 @@ def test_subnet_laplace_init(model, likelihood):
     subnetmask.select()
 
     # subnet Laplace with full Hessian should work
-    hessian_structure = 'full'
-    lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                  subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+    hessian_structure = "full"
+    lap = Laplace(
+        model,
+        likelihood=likelihood,
+        subset_of_weights="subnetwork",
+        subnetwork_indices=subnetmask.indices,
+        hessian_structure=hessian_structure,
+    )
     assert isinstance(lap, FullSubnetLaplace)
 
     # subnet Laplace with diagonal Hessian should work
-    hessian_structure = 'diag'
-    lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                  subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+    hessian_structure = "diag"
+    lap = Laplace(
+        model,
+        likelihood=likelihood,
+        subset_of_weights="subnetwork",
+        subnetwork_indices=subnetmask.indices,
+        hessian_structure=hessian_structure,
+    )
     assert isinstance(lap, DiagSubnetLaplace)
 
     # subnet Laplace without specifying subnetwork indices should raise an error
     with pytest.raises(TypeError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            hessian_structure=hessian_structure,
+        )
 
     # subnet Laplace with kron or lowrank Hessians should raise errors
-    hessian_structure = 'kron'
+    hessian_structure = "kron"
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
-    hessian_structure = 'lowrank'
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetmask.indices,
+            hessian_structure=hessian_structure,
+        )
+    hessian_structure = "lowrank"
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
-
-
-@pytest.mark.parametrize('likelihood,hessian_structure', product(likelihoods, hessian_structures))
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetmask.indices,
+            hessian_structure=hessian_structure,
+        )
+
+
+@pytest.mark.parametrize(
+    "likelihood,hessian_structure", product(likelihoods, hessian_structures)
+)
 def test_subnet_laplace_large_init(large_model, likelihood, hessian_structure):
     # use random subnet mask for this test
     subnetwork_mask = RandomSubnetMask
@@ -97,10 +133,15 @@ def test_subnet_laplace_large_init(large_model, likelihood, hessian_structure):
     subnetmask = subnetwork_mask(**subnetmask_kwargs)
     subnetmask.select()
 
-    lap = Laplace(large_model, likelihood=likelihood, subset_of_weights='subnetwork',
-                  subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+    lap = Laplace(
+        large_model,
+        likelihood=likelihood,
+        subset_of_weights="subnetwork",
+        subnetwork_indices=subnetmask.indices,
+        hessian_structure=hessian_structure,
+    )
     assert lap.n_params_subnet == n_param_subnet
-    if hessian_structure == 'full':
+    if hessian_structure == "full":
         assert lap.H.shape == (lap.n_params_subnet, lap.n_params_subnet)
     else:
         assert lap.H.shape == (lap.n_params_subnet,)
@@ -109,118 +150,206 @@ def test_subnet_laplace_large_init(large_model, likelihood, hessian_structure):
     assert torch.allclose(H, lap.H)
 
 
-@pytest.mark.parametrize('likelihood,hessian_structure', product(likelihoods, hessian_structures))
-def test_custom_subnetwork_indices(model, likelihood, class_loader, reg_loader, hessian_structure):
-    loader = class_loader if likelihood == 'classification' else reg_loader
+@pytest.mark.parametrize(
+    "likelihood,hessian_structure", product(likelihoods, hessian_structures)
+)
+def test_custom_subnetwork_indices(
+    model, likelihood, class_loader, reg_loader, hessian_structure
+):
+    loader = class_loader if likelihood == "classification" else reg_loader
 
     # subnetwork indices that are None should raise an error
     subnetwork_indices = None
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are not PyTorch tensors should raise an error
     subnetwork_indices = [0, 5, 11, 42]
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are empty tensors should raise an error
     subnetwork_indices = torch.LongTensor([])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are scalar tensors should raise an error
     subnetwork_indices = torch.LongTensor(11)
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are not 1D PyTorch tensors should raise an error
     subnetwork_indices = torch.LongTensor([[0, 5], [11, 42]])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are double tensors should raise an error
     subnetwork_indices = torch.DoubleTensor([0.0, 5.0, 11.0, 42.0])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are float tensors should raise an error
     subnetwork_indices = torch.FloatTensor([0.0, 5.0, 11.0, 42.0])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are half tensors should raise an error
     subnetwork_indices = torch.HalfTensor([0.0, 5.0, 11.0, 42.0])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are int tensors should raise an error
     subnetwork_indices = torch.IntTensor([0, 5, 11, 42])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are short tensors should raise an error
     subnetwork_indices = torch.ShortTensor([0, 5, 11, 42])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are char tensors should raise an error
     subnetwork_indices = torch.CharTensor([0, 5, 11, 42])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that are bool tensors should raise an error
     subnetwork_indices = torch.BoolTensor([0, 5, 11, 42])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that contain elements smaller than zero should raise an error
     subnetwork_indices = torch.LongTensor([0, -1, -11])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that contain elements larger than n_params should raise an error
     subnetwork_indices = torch.LongTensor([model.n_params + 1, model.n_params + 42])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # subnetwork indices that contain duplicate entries should raise an error
     subnetwork_indices = torch.LongTensor([0, 0, 5, 11, 11, 42])
     with pytest.raises(ValueError):
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetwork_indices,
+            hessian_structure=hessian_structure,
+        )
 
     # Non-empty, 1-dimensional torch.LongTensor with valid entries should work
     subnetwork_indices = torch.LongTensor([0, 5, 11, 42])
-    lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                  subnetwork_indices=subnetwork_indices, hessian_structure=hessian_structure)
+    lap = Laplace(
+        model,
+        likelihood=likelihood,
+        subset_of_weights="subnetwork",
+        subnetwork_indices=subnetwork_indices,
+        hessian_structure=hessian_structure,
+    )
     lap.fit(loader)
     assert isinstance(lap, SubnetLaplace)
     assert lap.n_params_subnet == 4
-    if hessian_structure == 'full':
+    if hessian_structure == "full":
         assert lap.H.shape == (4, 4)
     else:
         assert lap.H.shape == (4,)
     assert lap.backend.subnetwork_indices.equal(subnetwork_indices)
 
 
-@pytest.mark.parametrize('subnetwork_mask,likelihood,hessian_structure',
-                         product(score_based_subnet_masks, likelihoods, hessian_structures))
-def test_score_based_subnet_masks(model, likelihood, subnetwork_mask, class_loader, reg_loader, hessian_structure):
-    loader = class_loader if likelihood == 'classification' else reg_loader
+@pytest.mark.parametrize(
+    "subnetwork_mask,likelihood,hessian_structure",
+    product(score_based_subnet_masks, likelihoods, hessian_structures),
+)
+def test_score_based_subnet_masks(
+    model, likelihood, subnetwork_mask, class_loader, reg_loader, hessian_structure
+):
+    loader = class_loader if likelihood == "classification" else reg_loader
     model_params = parameters_to_vector(model.parameters())
 
     # set subnetwork mask arguments
@@ -266,8 +395,13 @@ def test_score_based_subnet_masks(model, likelihood, subnetwork_mask, class_load
         subnetmask.select(loader)
 
     # define valid subnet Laplace model
-    lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                  subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+    lap = Laplace(
+        model,
+        likelihood=likelihood,
+        subset_of_weights="subnetwork",
+        subnetwork_indices=subnetmask.indices,
+        hessian_structure=hessian_structure,
+    )
     assert isinstance(lap, SubnetLaplace)
 
     # fit Laplace model
@@ -280,22 +414,30 @@ def test_score_based_subnet_masks(model, likelihood, subnetwork_mask, class_load
     assert parameters_to_vector(model.parameters()).equal(model_params)
 
     # check that Hessian and prior precision is of correct shape
-    if hessian_structure == 'full':
+    if hessian_structure == "full":
         assert lap.H.shape == (n_params_subnet, n_params_subnet)
     else:
         assert lap.H.shape == (n_params_subnet,)
     assert lap.prior_precision_diag.shape == (n_params_subnet,)
 
 
-@pytest.mark.parametrize('subnetwork_mask,likelihood,hessian_structure',
-                         product(layer_subnet_masks, likelihoods, hessian_structures))
-def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, reg_loader, hessian_structure):
-    loader = class_loader if likelihood == 'classification' else reg_loader
+@pytest.mark.parametrize(
+    "subnetwork_mask,likelihood,hessian_structure",
+    product(layer_subnet_masks, likelihoods, hessian_structures),
+)
+def test_layer_subnet_masks(
+    model, likelihood, subnetwork_mask, class_loader, reg_loader, hessian_structure
+):
+    loader = class_loader if likelihood == "classification" else reg_loader
     subnetmask_kwargs = dict(model=model)
 
     # fit last-layer Laplace model
-    lllap = Laplace(model, likelihood=likelihood, subset_of_weights='last_layer',
-                    hessian_structure=hessian_structure)
+    lllap = Laplace(
+        model,
+        likelihood=likelihood,
+        subset_of_weights="last_layer",
+        hessian_structure=hessian_structure,
+    )
     lllap.fit(loader)
 
     # should raise error if we pass number of subnet parameters
@@ -319,23 +461,28 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
             subnetmask.select(loader)
 
         # should raise error if we pass a parameter name list with invalid parameter names
-        subnetmask_kwargs.update(parameter_names=['123'])
+        subnetmask_kwargs.update(parameter_names=["123"])
         with pytest.raises(ValueError):
             subnetmask = subnetwork_mask(**subnetmask_kwargs)
             subnetmask.select(loader)
 
         # define last-layer Laplace model by parameter names and check that
         # Hessian is identical to that of a full LLLaplace model
-        subnetmask_kwargs.update(parameter_names=['1.weight', '1.bias'])
+        subnetmask_kwargs.update(parameter_names=["1.weight", "1.bias"])
         subnetmask = subnetwork_mask(**subnetmask_kwargs)
         subnetmask.select(loader)
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetmask.indices,
+            hessian_structure=hessian_structure,
+        )
         lap.fit(loader)
         assert torch.allclose(lllap.H, lap.H, rtol=1e-3)
 
         # define valid parameter name subnet mask
-        subnetmask_kwargs.update(parameter_names=['0.weight', '1.bias'])
+        subnetmask_kwargs.update(parameter_names=["0.weight", "1.bias"])
         subnetmask = subnetwork_mask(**subnetmask_kwargs)
 
         # should raise error if we access number of subnet parameters before selecting the subnet
@@ -345,8 +492,13 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
 
         # select subnet mask and fit Laplace model
         subnetmask.select(loader)
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetmask.indices,
+            hessian_structure=hessian_structure,
+        )
         lap.fit(loader)
         assert isinstance(lap, SubnetLaplace)
 
@@ -364,23 +516,28 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
             subnetmask.select(loader)
 
         # should raise error if we pass a module name list with invalid module names
-        subnetmask_kwargs.update(module_names=['123'])
+        subnetmask_kwargs.update(module_names=["123"])
         with pytest.raises(ValueError):
             subnetmask = subnetwork_mask(**subnetmask_kwargs)
             subnetmask.select(loader)
 
         # define last-layer Laplace model by module name and check that
         # Hessian is identical to that of a full LLLaplace model
-        subnetmask_kwargs.update(module_names=['1'])
+        subnetmask_kwargs.update(module_names=["1"])
         subnetmask = subnetwork_mask(**subnetmask_kwargs)
         subnetmask.select(loader)
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetmask.indices,
+            hessian_structure=hessian_structure,
+        )
         lap.fit(loader)
         assert torch.allclose(lllap.H, lap.H, rtol=1e-3)
 
         # define valid parameter name subnet mask
-        subnetmask_kwargs.update(module_names=['0'])
+        subnetmask_kwargs.update(module_names=["0"])
         subnetmask = subnetwork_mask(**subnetmask_kwargs)
 
         # should raise error if we access number of subnet parameters before selecting the subnet
@@ -390,14 +547,19 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
 
         # select subnet mask and fit Laplace model
         subnetmask.select(loader)
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetmask.indices,
+            hessian_structure=hessian_structure,
+        )
         lap.fit(loader)
         assert isinstance(lap, SubnetLaplace)
 
     elif subnetwork_mask == LastLayerSubnetMask:
         # should raise error if we pass invalid last-layer name
-        subnetmask_kwargs.update(last_layer_name='123')
+        subnetmask_kwargs.update(last_layer_name="123")
         with pytest.raises(KeyError):
             subnetmask = subnetwork_mask(**subnetmask_kwargs)
             subnetmask.select(loader)
@@ -412,8 +574,13 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
 
         # select subnet mask and fit Laplace model
         subnetmask.select(loader)
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetmask.indices,
+            hessian_structure=hessian_structure,
+        )
         lap.fit(loader)
         assert isinstance(lap, SubnetLaplace)
 
@@ -421,7 +588,7 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
         assert torch.allclose(lllap.H, lap.H, rtol=1e-3)
 
         # define valid last-layer subnet mask (with passing the last-layer name)
-        subnetmask_kwargs.update(last_layer_name='1')
+        subnetmask_kwargs.update(last_layer_name="1")
         subnetmask = subnetwork_mask(**subnetmask_kwargs)
 
         # should raise error if we access number of subnet parameters before selecting the subnet
@@ -431,8 +598,13 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
 
         # select subnet mask and fit Laplace model
         subnetmask.select(loader)
-        lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                      subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+        lap = Laplace(
+            model,
+            likelihood=likelihood,
+            subset_of_weights="subnetwork",
+            subnetwork_indices=subnetmask.indices,
+            hessian_structure=hessian_structure,
+        )
         lap.fit(loader)
         assert isinstance(lap, SubnetLaplace)
 
@@ -445,16 +617,20 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
     assert lap.n_params_subnet == n_params_subnet
 
     # check that Hessian and prior precision is of correct shape
-    if hessian_structure == 'full':
+    if hessian_structure == "full":
         assert lap.H.shape == (n_params_subnet, n_params_subnet)
     else:
         assert lap.H.shape == (n_params_subnet,)
     assert lap.prior_precision_diag.shape == (n_params_subnet,)
 
 
-@pytest.mark.parametrize('likelihood,hessian_structure', product(likelihoods, hessian_structures))
-def test_full_subnet_mask(model, likelihood, class_loader, reg_loader, hessian_structure):
-    loader = class_loader if likelihood == 'classification' else reg_loader
+@pytest.mark.parametrize(
+    "likelihood,hessian_structure", product(likelihoods, hessian_structures)
+)
+def test_full_subnet_mask(
+    model, likelihood, class_loader, reg_loader, hessian_structure
+):
+    loader = class_loader if likelihood == "classification" else reg_loader
 
     # define full model 'subnet' mask class (i.e. where all parameters are part of the subnet)
     class FullSubnetMask(SubnetMask):
@@ -465,8 +641,13 @@ def get_subnet_mask(self, train_loader):
     subnetwork_mask = FullSubnetMask
     subnetmask = subnetwork_mask(model=model)
     subnetmask.select(loader)
-    lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                  subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+    lap = Laplace(
+        model,
+        likelihood=likelihood,
+        subset_of_weights="subnetwork",
+        subnetwork_indices=subnetmask.indices,
+        hessian_structure=hessian_structure,
+    )
     lap.fit(loader)
     assert isinstance(lap, SubnetLaplace)
 
@@ -476,31 +657,42 @@ def get_subnet_mask(self, train_loader):
     assert lap.n_params_subnet == model.n_params
 
     # check that the Hessian is identical to that of an all-weights Laplace model
-    full_lap = Laplace(model, likelihood=likelihood, subset_of_weights='all',
-                       hessian_structure=hessian_structure)
+    full_lap = Laplace(
+        model,
+        likelihood=likelihood,
+        subset_of_weights="all",
+        hessian_structure=hessian_structure,
+    )
     full_lap.fit(loader)
     assert torch.allclose(full_lap.H, lap.H, rtol=1e-3)
 
 
-@pytest.mark.parametrize('subnetwork_mask,hessian_structure', product(all_subnet_masks, hessian_structures))
+@pytest.mark.parametrize(
+    "subnetwork_mask,hessian_structure", product(all_subnet_masks, hessian_structures)
+)
 def test_regression_predictive(model, reg_loader, subnetwork_mask, hessian_structure):
     subnetmask_kwargs = dict(model=model)
     if subnetwork_mask in score_based_subnet_masks:
         subnetmask_kwargs.update(n_params_subnet=32)
         if subnetwork_mask == LargestVarianceSWAGSubnetMask:
-            subnetmask_kwargs.update(likelihood='regression')
+            subnetmask_kwargs.update(likelihood="regression")
         elif subnetwork_mask == LargestVarianceDiagLaplaceSubnetMask:
-            diag_laplace_model = DiagLaplace(model, 'regression')
+            diag_laplace_model = DiagLaplace(model, "regression")
             subnetmask_kwargs.update(diag_laplace_model=diag_laplace_model)
     elif subnetwork_mask == ParamNameSubnetMask:
-        subnetmask_kwargs.update(parameter_names=['0.weight', '1.bias'])
+        subnetmask_kwargs.update(parameter_names=["0.weight", "1.bias"])
     elif subnetwork_mask == ModuleNameSubnetMask:
-        subnetmask_kwargs.update(module_names=['0'])
+        subnetmask_kwargs.update(module_names=["0"])
 
     subnetmask = subnetwork_mask(**subnetmask_kwargs)
     subnetmask.select(reg_loader)
-    lap = Laplace(model, likelihood='regression', subset_of_weights='subnetwork',
-                  subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+    lap = Laplace(
+        model,
+        likelihood="regression",
+        subset_of_weights="subnetwork",
+        subnetwork_indices=subnetmask.indices,
+        hessian_structure=hessian_structure,
+    )
     assert isinstance(lap, SubnetLaplace)
 
     lap.fit(reg_loader)
@@ -509,40 +701,49 @@ def test_regression_predictive(model, reg_loader, subnetwork_mask, hessian_struc
 
     # error
     with pytest.raises(ValueError):
-        lap(X, pred_type='linear')
+        lap(X, pred_type="linear")
 
     # GLM predictive
-    f_mu, f_var = lap(X, pred_type='glm')
+    f_mu, f_var = lap(X, pred_type="glm")
     assert torch.allclose(f_mu, f)
     assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1], f_mu.shape[1]])
     assert len(f_mu) == len(X)
 
     # NN predictive (only diagonal variance estimation)
-    f_mu, f_var = lap(X, pred_type='nn', link_approx='mc')
+    f_mu, f_var = lap(X, pred_type="nn", link_approx="mc")
     assert f_mu.shape == f_var.shape
     assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1]])
     assert len(f_mu) == len(X)
 
 
-@pytest.mark.parametrize('subnetwork_mask,hessian_structure', product(all_subnet_masks, hessian_structures))
-def test_classification_predictive(model, class_loader, subnetwork_mask, hessian_structure):
+@pytest.mark.parametrize(
+    "subnetwork_mask,hessian_structure", product(all_subnet_masks, hessian_structures)
+)
+def test_classification_predictive(
+    model, class_loader, subnetwork_mask, hessian_structure
+):
     subnetmask_kwargs = dict(model=model)
     if subnetwork_mask in score_based_subnet_masks:
         subnetmask_kwargs.update(n_params_subnet=32)
         if subnetwork_mask == LargestVarianceSWAGSubnetMask:
-            subnetmask_kwargs.update(likelihood='classification')
+            subnetmask_kwargs.update(likelihood="classification")
         elif subnetwork_mask == LargestVarianceDiagLaplaceSubnetMask:
-            diag_laplace_model = DiagLaplace(model, 'classification')
+            diag_laplace_model = DiagLaplace(model, "classification")
             subnetmask_kwargs.update(diag_laplace_model=diag_laplace_model)
     elif subnetwork_mask == ParamNameSubnetMask:
-        subnetmask_kwargs.update(parameter_names=['0.weight', '1.bias'])
+        subnetmask_kwargs.update(parameter_names=["0.weight", "1.bias"])
     elif subnetwork_mask == ModuleNameSubnetMask:
-        subnetmask_kwargs.update(module_names=['0'])
+        subnetmask_kwargs.update(module_names=["0"])
 
     subnetmask = subnetwork_mask(**subnetmask_kwargs)
     subnetmask.select(class_loader)
-    lap = Laplace(model, likelihood='classification', subset_of_weights='subnetwork',
-                  subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+    lap = Laplace(
+        model,
+        likelihood="classification",
+        subset_of_weights="subnetwork",
+        subnetwork_indices=subnetmask.indices,
+        hessian_structure=hessian_structure,
+    )
     assert isinstance(lap, SubnetLaplace)
 
     lap.fit(class_loader)
@@ -551,31 +752,45 @@ def test_classification_predictive(model, class_loader, subnetwork_mask, hessian
 
     # error
     with pytest.raises(ValueError):
-        lap(X, pred_type='linear')
+        lap(X, pred_type="linear")
 
     # GLM predictive
-    f_pred = lap(X, pred_type='glm', link_approx='mc', n_samples=100)
+    f_pred = lap(X, pred_type="glm", link_approx="mc", n_samples=100)
     assert f_pred.shape == f.shape
-    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
-    f_pred = lap(X, pred_type='glm', link_approx='probit')
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
+    f_pred = lap(X, pred_type="glm", link_approx="probit")
     assert f_pred.shape == f.shape
-    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
-    f_pred = lap(X, pred_type='glm', link_approx='bridge')
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
+    f_pred = lap(X, pred_type="glm", link_approx="bridge")
     assert f_pred.shape == f.shape
-    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
-    f_pred = lap(X, pred_type='glm', link_approx='bridge_norm')
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
+    f_pred = lap(X, pred_type="glm", link_approx="bridge_norm")
     assert f_pred.shape == f.shape
-    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
 
     # NN predictive
-    f_pred = lap(X, pred_type='nn', link_approx='mc', n_samples=100)
+    f_pred = lap(X, pred_type="nn", link_approx="mc", n_samples=100)
     assert f_pred.shape == f.shape
-    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
-
-
-@pytest.mark.parametrize('subnetwork_mask,likelihood,hessian_structure',
-                         product(all_subnet_masks, likelihoods, hessian_structures))
-def test_subnet_marginal_likelihood(model, subnetwork_mask, likelihood, hessian_structure, class_loader, reg_loader):
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
+
+
+@pytest.mark.parametrize(
+    "subnetwork_mask,likelihood,hessian_structure",
+    product(all_subnet_masks, likelihoods, hessian_structures),
+)
+def test_subnet_marginal_likelihood(
+    model, subnetwork_mask, likelihood, hessian_structure, class_loader, reg_loader
+):
     subnetmask_kwargs = dict(model=model)
     if subnetwork_mask in score_based_subnet_masks:
         subnetmask_kwargs.update(n_params_subnet=32)
@@ -585,15 +800,20 @@ def test_subnet_marginal_likelihood(model, subnetwork_mask, likelihood, hessian_
             diag_laplace_model = DiagLaplace(model, likelihood)
             subnetmask_kwargs.update(diag_laplace_model=diag_laplace_model)
     elif subnetwork_mask == ParamNameSubnetMask:
-        subnetmask_kwargs.update(parameter_names=['0.weight', '1.bias'])
+        subnetmask_kwargs.update(parameter_names=["0.weight", "1.bias"])
     elif subnetwork_mask == ModuleNameSubnetMask:
-        subnetmask_kwargs.update(module_names=['0'])
+        subnetmask_kwargs.update(module_names=["0"])
 
     subnetmask = subnetwork_mask(**subnetmask_kwargs)
-    loader = class_loader if likelihood == 'classification' else reg_loader
+    loader = class_loader if likelihood == "classification" else reg_loader
     subnetmask.select(loader)
-    lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
-                  subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
+    lap = Laplace(
+        model,
+        likelihood=likelihood,
+        subset_of_weights="subnetwork",
+        subnetwork_indices=subnetmask.indices,
+        hessian_structure=hessian_structure,
+    )
     assert isinstance(lap, SubnetLaplace)
 
     lap.fit(loader)
diff --git a/tests/test_subset_params.py b/tests/test_subset_params.py
index 0105327..21ddb4a 100644
--- a/tests/test_subset_params.py
+++ b/tests/test_subset_params.py
@@ -1,21 +1,14 @@
-from math import sqrt
-import pytest
 from itertools import product
-import numpy as np
-from copy import deepcopy
+
+import pytest
 import torch
 from torch import nn
-from torch.distributions.multivariate_normal import MultivariateNormal
 from torch.nn.utils import parameters_to_vector
 from torch.utils.data import DataLoader, TensorDataset
-from torch.distributions import Normal, Categorical
-from torchvision.models import wide_resnet50_2
-
-from laplace.laplace import FullLaplace, KronLaplace, DiagLaplace, LowRankLaplace
-from laplace.utils import KronDecomposed
-from laplace.curvature import AsdlGGN, AsdlHessian, AsdlEF, BackPackEF, BackPackGGN
-from tests.utils import jacobians_naive
 
+from laplace.curvature import AsdlEF, AsdlGGN, AsdlHessian, BackPackEF, BackPackGGN
+from laplace.curvature.curvlinops import CurvlinopsEF, CurvlinopsGGN, CurvlinopsHessian
+from laplace.laplace import DiagLaplace, FullLaplace, KronLaplace
 
 torch.manual_seed(240)
 torch.set_default_tensor_type(torch.DoubleTensor)
@@ -25,10 +18,10 @@
 @pytest.fixture
 def model():
     model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
-    setattr(model, 'output_size', 2)
+    setattr(model, "output_size", 2)
     model_params = list(model.parameters())
-    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
-    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    setattr(model, "n_layers", len(model_params))  # number of parameter groups
+    setattr(model, "n_params", len(parameters_to_vector(model_params)))
 
     # Subset of params
     for p in model.parameters():
@@ -52,69 +45,77 @@ def reg_loader():
     return DataLoader(TensorDataset(X, y), batch_size=3)
 
 
-@pytest.mark.parametrize('laplace,lh', product(flavors, ['classification', 'regression']))
-def test_incompatible_backend(laplace, lh, model):
-    lap = laplace(model, lh, backend=AsdlEF)
-    lap = laplace(model, lh, backend=AsdlGGN)
-    lap = laplace(model, lh, backend=AsdlHessian)
+@pytest.mark.parametrize(
+    "laplace,lh", product(flavors, ["classification", "regression"])
+)
+def test_compatible_backend(laplace, lh, model):
+    laplace(model, lh, backend=CurvlinopsEF)
+    laplace(model, lh, backend=CurvlinopsGGN)
+    laplace(model, lh, backend=CurvlinopsHessian)
+    laplace(model, lh, backend=AsdlEF)
+    laplace(model, lh, backend=AsdlGGN)
+    laplace(model, lh, backend=AsdlHessian)
 
 
-@pytest.mark.parametrize('laplace,lh', product(flavors, ['classification', 'regression']))
+@pytest.mark.parametrize(
+    "laplace,lh", product(flavors, ["classification", "regression"])
+)
 def test_incompatible_backend(laplace, lh, model):
     with pytest.raises(ValueError):
-        lap = laplace(model, lh, backend=BackPackGGN)
+        laplace(model, lh, backend=BackPackGGN)
 
     with pytest.raises(ValueError):
-        lap = laplace(model, lh, backend=BackPackEF)
+        laplace(model, lh, backend=BackPackEF)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_mean_clf(laplace, model, class_loader):
     n_params = model[0].weight.numel()
-    lap = laplace(model, 'classification')
+    lap = laplace(model, "classification")
     lap.fit(class_loader)
     assert lap.mean.shape == (n_params,)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_mean_reg(laplace, model, reg_loader):
     n_params = model[0].weight.numel()
-    lap = laplace(model, 'regression')
+    lap = laplace(model, "regression")
     lap.fit(reg_loader)
     assert lap.mean.shape == (n_params,)
 
 
 def test_post_precision_diag(model, class_loader):
     n_params = model[0].weight.numel()
-    lap = DiagLaplace(model, 'classification')
+    lap = DiagLaplace(model, "classification")
     lap.fit(class_loader)
     assert lap.posterior_precision.shape == (n_params,)
 
 
 def test_post_precision_kron(model, class_loader):
     n_params = model[0].weight.numel()
-    lap = KronLaplace(model, 'classification')
+    lap = KronLaplace(model, "classification")
     lap.fit(class_loader)
     assert lap.posterior_precision.to_matrix().shape == (n_params, n_params)
 
 
-@pytest.mark.parametrize('laplace', flavors)
+@pytest.mark.parametrize("laplace", flavors)
 def test_predictive(laplace, model, class_loader):
-    lap = laplace(model, 'classification')
+    lap = laplace(model, "classification")
     lap.fit(class_loader)
-    lap(torch.randn(5, 3), pred_type='nn', link_approx='mc')
+    lap(torch.randn(5, 3), pred_type="nn", link_approx="mc")
 
 
-@pytest.mark.parametrize('laplace', flavors)
-def test_marglik(laplace, model, class_loader):
-    lap = laplace(model, 'classification')
+@pytest.mark.parametrize("laplace", flavors)
+def test_marglik_glm(laplace, model, class_loader):
+    lap = laplace(model, "classification")
     lap.fit(class_loader)
-    lap.optimize_prior_precision(method='marglik')
+    lap.optimize_prior_precision(method="marglik")
 
 
-@pytest.mark.parametrize('laplace', flavors)
-def test_marglik(laplace, model, class_loader):
-    lap = laplace(model, 'classification')
+@pytest.mark.parametrize("laplace", flavors)
+def test_marglik_nn(laplace, model, class_loader):
+    lap = laplace(model, "classification")
     lap.fit(class_loader)
-    lap.optimize_prior_precision(method='gridsearch', val_loader=class_loader,
-                                 pred_type='nn', link_approx='mc')
+    lap.optimize_prior_precision(
+        method="gridsearch", val_loader=class_loader, pred_type="nn", link_approx="mc"
+    )
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 29c142b..36b5bd1 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,16 +1,18 @@
+import math
+
 import torch
-from torch.utils.data import TensorDataset, DataLoader
+from torch.utils.data import DataLoader, TensorDataset
+
 from laplace import Laplace
 from laplace.utils import (
-    invsqrt_precision,
+    RunningNLLMetric,
     diagonal_add_scalar,
-    symeig,
+    get_nll,
+    invsqrt_precision,
     normal_samples,
+    symeig,
     validate,
-    get_nll,
-    RunningNLLMetric,
 )
-import math
 
 
 def test_sqrt_precision():
@@ -30,7 +32,7 @@ def test_diagonal_add_scalar():
 def test_symeig_custom():
     X = torch.randn(20, 100)
     M = X @ X.T
-    l1, W1 = torch.linalg.eigh(M, UPLO='U')
+    l1, W1 = torch.linalg.eigh(M, UPLO="U")
     l2, W2 = symeig(M)
     assert torch.allclose(l1, l2)
     assert torch.allclose(W1, W2)
@@ -39,7 +41,7 @@ def test_symeig_custom():
 def test_symeig_custom_low_rank():
     X = torch.randn(1000, 10)
     M = X @ X.T
-    l1, W1 = torch.linalg.eigh(M, UPLO='U')
+    l1, W1 = torch.linalg.eigh(M, UPLO="U")
     l2, W2 = symeig(M)
     # symeig should fail for low-rank
     assert not torch.all(l1 >= 0.0)
@@ -82,11 +84,11 @@ def test_validate():
     model = torch.nn.Sequential(
         torch.nn.Linear(10, 20), torch.nn.ReLU(), torch.nn.Linear(20, 3)
     )
-    la = Laplace(model, 'classification', 'all')
+    la = Laplace(model, "classification", "all")
     la.fit(dataloader)
 
     res = validate(
-        la, dataloader, get_nll, pred_type='nn', link_approx='mc', n_samples=10
+        la, dataloader, get_nll, pred_type="nn", link_approx="mc", n_samples=10
     )
     assert res != math.nan
     assert isinstance(res, float)
@@ -96,8 +98,8 @@ def test_validate():
         la,
         dataloader,
         RunningNLLMetric(),
-        pred_type='nn',
-        link_approx='mc',
+        pred_type="nn",
+        link_approx="mc",
         n_samples=10,
     )
     assert res != math.nan
diff --git a/tests/utils.py b/tests/utils.py
index c9e136b..0e58b57 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,6 +1,7 @@
 from collections import UserDict
 from collections.abc import MutableMapping
 from typing import Any, List
+
 import torch
 from torch.utils.data import Dataset