diff --git a/README.md b/README.md
index f5b7035..1a788d9 100644
--- a/README.md
+++ b/README.md
@@ -86,18 +86,18 @@ class. It's
 [**init**](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.trainer.trainer.html#pytorch_lightning.trainer.trainer.Trainer.__init__)
 method provides various configuration options.
 
-If you want to run K-Means with a GPU, you can pass the option `gpus=1` to the estimator's
-initializer:
+If you want to run K-Means with a GPU, you can pass the options `accelerator='gpu'` and `devices=1`
+to the estimator's initializer:
 
 ```python
-estimator = KMeans(3, trainer_params=dict(gpus=1))
+estimator = KMeans(3, trainer_params=dict(accelerator='gpu', devices=1))
 ```
 
 Similarly, if you want to train on 4 nodes simultaneously where each node has one GPU available,
 you can specify this as follows:
 
 ```python
-estimator = KMeans(3, trainer_params=dict(num_nodes=4, gpus=1))
+estimator = KMeans(3, trainer_params=dict(num_nodes=4, accelerator='gpu', devices=1))
 ```
 
 In fact, **you do not need to change anything else in your code**.
diff --git a/docs/index.rst b/docs/index.rst
index f924d3f..64697e6 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -78,19 +78,19 @@ For GPU- and multi-node training, PyCave leverages PyTorch Lightning. The hardwa
 runs on is determined by the :class:`pytorch_lightning.trainer.Trainer` class. It's
 :meth:`~pytorch_lightning.trainer.Trainer.__init__` method provides various configuration options.
 
-If you want to run K-Means with a GPU, you can pass the option ``gpus=1`` to the estimator's
+If you want to run K-Means with a GPU, you can pass the option ``accelerator='gpu'`` and ``devices=1`` to the estimator's
 initializer:
 
 .. code-block:: python
 
-    estimator = KMeans(3, trainer_params=dict(gpus=1))
+    estimator = KMeans(3, trainer_params=dict(accelerator='gpu', devices=1))
 
 Similarly, if you want to train on 4 nodes simultaneously where each node has one GPU available,
 you can specify this as follows:
 
 .. code-block:: python
 
-    estimator = KMeans(3, trainer_params=dict(num_nodes=4, gpus=1))
+    estimator = KMeans(3, trainer_params=dict(num_nodes=4, accelerator='gpu', 1))
 
 In fact, **you do not need to change anything else in your code**.
 
diff --git a/poetry.lock b/poetry.lock
index d53d8f3..c7f021f 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -681,14 +681,14 @@ tomli = ["tomli (<2.0.0)"]
 
 [[package]]
 name = "docutils"
-version = "0.17.1"
+version = "0.19"
 description = "Docutils -- Python Documentation Utilities"
 category = "dev"
 optional = false
-python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
+python-versions = ">=3.7"
 files = [
-    {file = "docutils-0.17.1-py2.py3-none-any.whl", hash = "sha256:cf316c8370a737a022b72b56874f6602acf974a37a9fba42ec2876387549fc61"},
-    {file = "docutils-0.17.1.tar.gz", hash = "sha256:686577d2e4c32380bb50cbb22f575ed742d58168cee37e99117a854bcd88f125"},
+    {file = "docutils-0.19-py3-none-any.whl", hash = "sha256:5e1de4d849fee02c63b040a4a3fd567f4ab104defd8a5511fbbc24a8a017efbc"},
+    {file = "docutils-0.19.tar.gz", hash = "sha256:33995a6753c30b7f577febfc2c50411fec6aac7f7ffeb7c4cfe5991072dcf9e6"},
 ]
 
 [[package]]
@@ -3157,22 +3157,23 @@ test = ["cython", "html5lib", "pytest (>=4.6)", "typed_ast"]
 
 [[package]]
 name = "sphinx-autodoc-typehints"
-version = "1.19.1"
+version = "1.19.5"
 description = "Type hints (PEP 484) support for the Sphinx autodoc extension"
 category = "dev"
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "sphinx_autodoc_typehints-1.19.1-py3-none-any.whl", hash = "sha256:9be46aeeb1b315eb5df1f3a7cb262149895d16c7d7dcd77b92513c3c3a1e85e6"},
-    {file = "sphinx_autodoc_typehints-1.19.1.tar.gz", hash = "sha256:6c841db55e0e9be0483ff3962a2152b60e79306f4288d8c4e7e86ac84486a5ea"},
+    {file = "sphinx_autodoc_typehints-1.19.5-py3-none-any.whl", hash = "sha256:ea55b3cc3f485e3a53668bcdd08de78121ab759f9724392fdb5bf3483d786328"},
+    {file = "sphinx_autodoc_typehints-1.19.5.tar.gz", hash = "sha256:38a227378e2bc15c84e29af8cb1d7581182da1107111fd1c88b19b5eb7076205"},
 ]
 
 [package.dependencies]
-Sphinx = ">=4.5"
+sphinx = ">=5.3"
 
 [package.extras]
-testing = ["covdefaults (>=2.2)", "coverage (>=6.3)", "diff-cover (>=6.4)", "nptyping (>=2.1.2)", "pytest (>=7.1)", "pytest-cov (>=3)", "sphobjinv (>=2)", "typing-extensions (>=4.1)"]
-type-comments = ["typed-ast (>=1.5.2)"]
+docs = ["furo (>=2022.9.29)", "sphinx (>=5.3)", "sphinx-autodoc-typehints (>=1.19.4)"]
+testing = ["covdefaults (>=2.2)", "coverage (>=6.5)", "diff-cover (>=7.0.1)", "nptyping (>=2.3.1)", "pytest (>=7.2)", "pytest-cov (>=4)", "sphobjinv (>=2.2.2)", "typing-extensions (>=4.4)"]
+type-comment = ["typed-ast (>=1.5.4)"]
 
 [[package]]
 name = "sphinx-automodapi"
@@ -3529,14 +3530,14 @@ typing-extensions = "*"
 
 [[package]]
 name = "torchmetrics"
-version = "0.5.1"
+version = "0.6.2"
 description = "PyTorch native Metrics"
 category = "main"
 optional = false
 python-versions = ">=3.6"
 files = [
-    {file = "torchmetrics-0.5.1-py3-none-any.whl", hash = "sha256:4e5497bc5c9d19fa520748cda89f6d863868bb5be33ec47d2834c0988bf737c5"},
-    {file = "torchmetrics-0.5.1.tar.gz", hash = "sha256:22fbcb6fc05348ca3f2bd06e0763e88411a6b68c2b9fc26084b39d40cc4021b0"},
+    {file = "torchmetrics-0.6.2-py3-none-any.whl", hash = "sha256:2b8a75e7af97bbe2308a00823ec90a222c83f1e4f098e4020c1c863c81869412"},
+    {file = "torchmetrics-0.6.2.tar.gz", hash = "sha256:a43e232a73af71febce98949122600bfafb9437eefbda28fb36f328bc02357fb"},
 ]
 
 [package.dependencies]
@@ -3545,9 +3546,14 @@ packaging = "*"
 torch = ">=1.3.1"
 
 [package.extras]
-all = ["nltk (>=3.6)", "scipy", "torch-fidelity", "torchvision", "tqdm (>=4.41.0)"]
-image = ["scipy", "torch-fidelity", "torchvision"]
-text = ["nltk (>=3.6)", "tqdm (>=4.41.0)"]
+all = ["bert-score (==0.3.10)", "check-manifest", "cloudpickle (>=1.3)", "codecov (>=2.1)", "coverage (>5.2)", "docutils (>=0.16)", "flake8", "isort (>=5.0)", "jiwer (>=2.3.0)", "lpips", "mir-eval (>=0.6)", "mypy (>=0.790)", "myst-parser", "nbsphinx (>=0.8)", "nltk (>=3.6)", "pandoc (>=1.0)", "pesq (>=0.0.3)", "phmdoctest (>=1.1.1)", "pre-commit (>=1.0)", "pypesq", "pystoi", "pytest (>=6.0)", "pytest-cov (>2.10)", "pytorch-lightning (>=1.0)", "pytorch-lightning (>=1.1)", "regex (>=2021.9.24)", "rouge-score (>=0.0.4)", "sacrebleu (>=2.0.0)", "scikit-image (>0.17.1)", "scikit-learn (>=0.24)", "scipy", "sphinx (>=4.0)", "sphinx-autodoc-typehints (>=1.0)", "sphinx-copybutton (>=0.3)", "sphinx-paramlinks (>=0.5.1)", "sphinx-togglebutton (>=0.2)", "sphinxcontrib-fulltoc (>=1.0)", "sphinxcontrib-mockautodoc", "torch-fidelity", "torchvision", "tqdm (>=4.41.0)", "transformers (>=4.0)", "twine (>=3.2)"]
+audio = ["pesq (>=0.0.3)", "pystoi"]
+detection = ["torchvision"]
+docs = ["docutils (>=0.16)", "myst-parser", "nbsphinx (>=0.8)", "pandoc (>=1.0)", "pytorch-lightning (>=1.1)", "sphinx (>=4.0)", "sphinx-autodoc-typehints (>=1.0)", "sphinx-copybutton (>=0.3)", "sphinx-paramlinks (>=0.5.1)", "sphinx-togglebutton (>=0.2)", "sphinxcontrib-fulltoc (>=1.0)", "sphinxcontrib-mockautodoc"]
+image = ["lpips", "scipy", "torch-fidelity", "torchvision"]
+integrate = ["pytorch-lightning (>=1.0)"]
+test = ["bert-score (==0.3.10)", "check-manifest", "cloudpickle (>=1.3)", "codecov (>=2.1)", "coverage (>5.2)", "flake8", "isort (>=5.0)", "jiwer (>=2.3.0)", "mir-eval (>=0.6)", "mypy (>=0.790)", "phmdoctest (>=1.1.1)", "pre-commit (>=1.0)", "pypesq", "pytest (>=6.0)", "pytest-cov (>2.10)", "rouge-score (>=0.0.4)", "sacrebleu (>=2.0.0)", "scikit-image (>0.17.1)", "scikit-learn (>=0.24)", "transformers (>=4.0)", "twine (>=3.2)"]
+text = ["nltk (>=3.6)", "regex (>=2021.9.24)", "tqdm (>=4.41.0)"]
 
 [[package]]
 name = "tornado"
@@ -3943,4 +3949,4 @@ testing = ["flake8 (<5)", "func-timeout", "jaraco.functools", "jaraco.itertools"
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8,<3.11"
-content-hash = "a0b9f839c9229558975cfdb29ba1796439fcf5efc1b1deb1d74a6c1bb70d7707"
+content-hash = "1f7761f930fcfbee9422d29031b81820ab79236a7e4bd7552eebb9e3e6adb828"
diff --git a/pycave/bayes/gmm/lightning_module.py b/pycave/bayes/gmm/lightning_module.py
index ab4073d..a5aaf90 100644
--- a/pycave/bayes/gmm/lightning_module.py
+++ b/pycave/bayes/gmm/lightning_module.py
@@ -2,7 +2,7 @@
 import pytorch_lightning as pl
 import torch
 from pytorch_lightning.callbacks import EarlyStopping
-from torchmetrics import AverageMeter
+from torchmetrics import MeanMetric
 from pycave.bayes.core import cholesky_precision
 from pycave.utils import NonparametricLightningModule
 from .metrics import CovarianceAggregator, MeanAggregator, PriorAggregator
@@ -65,7 +65,7 @@ def __init__(
         )
 
         # Initialize metrics
-        self.metric_nll = AverageMeter(dist_sync_fn=self.all_gather)
+        self.metric_nll = MeanMetric(dist_sync_fn=self.all_gather)
 
     def configure_callbacks(self) -> list[pl.Callback]:
         if self.convergence_tolerance == 0:
diff --git a/pycave/bayes/gmm/metrics.py b/pycave/bayes/gmm/metrics.py
index 94fe0dc..388639a 100644
--- a/pycave/bayes/gmm/metrics.py
+++ b/pycave/bayes/gmm/metrics.py
@@ -9,6 +9,8 @@ class PriorAggregator(Metric):
     The prior aggregator aggregates component probabilities over batches and process.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_components: int,
@@ -33,6 +35,8 @@ class MeanAggregator(Metric):
     The mean aggregator aggregates component means over batches and processes.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_components: int,
@@ -63,6 +67,8 @@ class CovarianceAggregator(Metric):
     The covariance aggregator aggregates component covariances over batches and processes.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_components: int,
diff --git a/pycave/bayes/markov_chain/lightning_module.py b/pycave/bayes/markov_chain/lightning_module.py
index c6f9074..80119bf 100644
--- a/pycave/bayes/markov_chain/lightning_module.py
+++ b/pycave/bayes/markov_chain/lightning_module.py
@@ -1,6 +1,6 @@
 import torch
 from torch.nn.utils.rnn import PackedSequence
-from torchmetrics import AverageMeter
+from torchmetrics import MeanMetric
 from pycave.bayes.markov_chain.metrics import StateCountAggregator
 from pycave.utils import NonparametricLightningModule
 from .model import MarkovChainModel
@@ -27,7 +27,7 @@ def __init__(self, model: MarkovChainModel, symmetric: bool = False):
             symmetric=self.symmetric,
             dist_sync_fn=self.all_gather,
         )
-        self.metric_nll = AverageMeter(dist_sync_fn=self.all_gather)
+        self.metric_nll = MeanMetric(dist_sync_fn=self.all_gather)
 
     def on_train_epoch_start(self) -> None:
         self.aggregator.reset()
diff --git a/pycave/bayes/markov_chain/metrics.py b/pycave/bayes/markov_chain/metrics.py
index 386306d..3b208ce 100644
--- a/pycave/bayes/markov_chain/metrics.py
+++ b/pycave/bayes/markov_chain/metrics.py
@@ -9,6 +9,8 @@ class StateCountAggregator(Metric):
     The state count aggregator aggregates initial states and transitions between states.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_states: int,
diff --git a/pycave/clustering/kmeans/lightning_module.py b/pycave/clustering/kmeans/lightning_module.py
index aba9793..78dd51e 100644
--- a/pycave/clustering/kmeans/lightning_module.py
+++ b/pycave/clustering/kmeans/lightning_module.py
@@ -4,7 +4,7 @@
 import pytorch_lightning as pl
 import torch
 from pytorch_lightning.callbacks import EarlyStopping
-from torchmetrics import AverageMeter
+from torchmetrics import MeanMetric
 from pycave.utils import NonparametricLightningModule
 from .metrics import (
     BatchAverager,
@@ -51,7 +51,7 @@ def __init__(
         )
 
         # Initialize metrics
-        self.metric_inertia = AverageMeter()
+        self.metric_inertia = MeanMetric()
 
     def configure_callbacks(self) -> List[pl.Callback]:
         if self.convergence_tolerance == 0:
@@ -239,8 +239,8 @@ def nonparametric_training_step(self, batch: torch.Tensor, batch_idx: int) -> No
 
     def nonparametric_training_epoch_end(self) -> None:
         if self.current_epoch == 0:
-            choice = self.uniform_sampler.compute()[0]
-            self.model.centroids[0].copy_(choice)
+            choice = self.uniform_sampler.compute()
+            self.model.centroids[0].copy_(choice[0] if choice.dim() > 0 else choice)
         elif self._is_current_epoch_sampling:
             candidates = self.distance_sampler.compute()
             self.centroid_candidates.copy_(candidates)
diff --git a/pycave/clustering/kmeans/metrics.py b/pycave/clustering/kmeans/metrics.py
index 2cb1cdd..1146f25 100644
--- a/pycave/clustering/kmeans/metrics.py
+++ b/pycave/clustering/kmeans/metrics.py
@@ -9,6 +9,8 @@ class CentroidAggregator(Metric):
     The centroid aggregator aggregates kmeans centroids over batches and processes.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_clusters: int,
@@ -49,6 +51,8 @@ class UniformSampler(Metric):
     they were already sampled from).
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_choices: int,
@@ -109,6 +113,8 @@ class DistanceSampler(Metric):
     duplicates.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_choices: int,
@@ -169,6 +175,8 @@ class BatchSummer(Metric):
     Sums the values for a batch of items independently.
     """
 
+    full_state_update = True
+
     def __init__(self, num_values: int, *, dist_sync_fn: Optional[Callable[[Any], Any]] = None):
         super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore
 
@@ -187,6 +195,8 @@ class BatchAverager(Metric):
     Averages the values for a batch of items independently.
     """
 
+    full_state_update = False
+
     def __init__(
         self,
         num_values: int,
diff --git a/pyproject.toml b/pyproject.toml
index 35b00dc..5b30d41 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,7 +18,7 @@ numpy = "^1.20.3"
 python = ">=3.8,<3.11"
 pytorch-lightning = "^1.6.0"
 torch = "^1.8.0"
-torchmetrics = "^0.5.1,<0.6.0"
+torchmetrics = "^0.6.0"
 
 [tool.poetry.group.pre-commit.dependencies]
 black = "^22.12.0"
diff --git a/tests/bayes/gmm/benchmark_gmm_estimator.py b/tests/bayes/gmm/benchmark_gmm_estimator.py
index 3ef025f..f4dc818 100644
--- a/tests/bayes/gmm/benchmark_gmm_estimator.py
+++ b/tests/bayes/gmm/benchmark_gmm_estimator.py
@@ -128,6 +128,6 @@ def test_pycave_gpu(
         convergence_tolerance=0,
         covariance_regularization=1e-3,
         batch_size=batch_size,
-        trainer_params=dict(max_epochs=100, gpus=1),
+        trainer_params=dict(max_epochs=100, accelerator="gpu", devices=1),
     )
     benchmark(estimator.fit, data)
diff --git a/tests/clustering/kmeans/benchmark_kmeans_estimator.py b/tests/clustering/kmeans/benchmark_kmeans_estimator.py
index 8aa418e..f95c399 100644
--- a/tests/clustering/kmeans/benchmark_kmeans_estimator.py
+++ b/tests/clustering/kmeans/benchmark_kmeans_estimator.py
@@ -120,6 +120,6 @@ def test_pycave_gpu(
         init_strategy=init_strategy,
         batch_size=batch_size,
         convergence_tolerance=0,
-        trainer_params=dict(gpus=1, max_epochs=100),
+        trainer_params=dict(max_epochs=100, accelerator="gpu", devices=1),
     )
     benchmark(estimator.fit, data)