Repository: datamol-io/graphium (language: python)

Issue:
Update `MetricWrapper` to work with the update/compute from torchmetrics
Currently, we cannot compute train metrics because it requires to save all the preds and targets over the epoch.
The problem also affects the val and test sets, but since they're smaller, the effect is less noticeable.

To fix this, the [MetricWrapper](https://graphium-docs.datamol.io/stable/api/graphium.trainer.html#metrics) should have an update method taking the preds and target that calls the underlying self.metrics.update, and the compute method which no longer takes in the pred and target, but instead calls the self.metric.compute

- [x] Add the `update` to the `MetricWrapper`
- [x] Modify the `MetricWrapper.compute` to work with the update
- [x] How to deal with missing labels??

Also, all [TorchMetrics in spaces.py](https://github.com/datamol-io/graphium/blob/41a1172a16d2da1138c87d4e472d36ba5f8d2e8b/graphium/utils/spaces.py#L81) should become their class equivalent rather than functions.

- [x] Change the `spaces.py` to use classes rather than functions. Make sure the classes get initialized.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 41dbc0e45..2afd2115e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -39,6 +39,7 @@ graphium/data/cache/
 graphium/data/b3lyp/
 graphium/data/PCQM4Mv2/
 graphium/data/PCQM4M/
+graphium/data/largemix/
 graphium/data/neurips2023/small-dataset/
 graphium/data/neurips2023/large-dataset/
 graphium/data/neurips2023/dummy-dataset/
diff --git a/docs/api/graphium.ipu.md b/docs/api/graphium.ipu.md
index 2fdf82416..3943e78a4 100644
--- a/docs/api/graphium.ipu.md
+++ b/docs/api/graphium.ipu.md
@@ -22,11 +22,6 @@ Code for adapting to run on IPU
 ::: graphium.ipu.ipu_losses
 
 
-## IPU Metrics
-------------
-::: graphium.ipu.ipu_metrics
-
-
 ## IPU Simple Lightning
 ------------
 ::: graphium.ipu.ipu_simple_lightning
diff --git a/env.yml b/env.yml
index 62bf95a3c..d3a3cbc23 100644
--- a/env.yml
+++ b/env.yml
@@ -31,7 +31,7 @@ dependencies:
   - cuda-version # works also with CPU-only system.
   - pytorch >=1.12
   - lightning >=2.0
-  - torchmetrics >=0.7.0,<0.11
+  - torchmetrics
   - ogb
   - pytorch_geometric >=2.0 # Use `pyg` for Windows instead of `pytorch_geometric`
   - wandb
diff --git a/expts/hydra-configs/tasks/loss_metrics_datamodule/toymix.yaml b/expts/hydra-configs/tasks/loss_metrics_datamodule/toymix.yaml
index 9ac744a52..ceb4e6d69 100644
--- a/expts/hydra-configs/tasks/loss_metrics_datamodule/toymix.yaml
+++ b/expts/hydra-configs/tasks/loss_metrics_datamodule/toymix.yaml
@@ -1,6 +1,8 @@
 # @package _global_
 
 predictor:
+  target_nan_mask: ignore
+  multitask_handling: flatten
   metrics_on_progress_bar:
     qm9: ["mae"]
     tox21: ["auroc"]
@@ -13,29 +15,31 @@ predictor:
 metrics:
   qm9: &qm9_metrics
     - name: mae
-      metric: mae_ipu
-      target_nan_mask: null
+      metric: mae
+      target_nan_mask: ignore
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
-      target_nan_mask: null
+      target_nan_mask: ignore
       multitask_handling: mean-per-label
     - name: r2_score
-      metric: r2_score_ipu
-      target_nan_mask: null
+      metric: r2_score
+      target_nan_mask: ignore
       multitask_handling: mean-per-label
       threshold_kwargs: null
   tox21:
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       task: binary
+      target_nan_mask: ignore
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: averageprecision
       task: binary
+      target_nan_mask: ignore
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: f1 > 0.5
@@ -44,6 +48,8 @@ metrics:
       target_to_int: True
       num_classes: 2
       average: micro
+      task: binary
+      target_nan_mask: ignore
       threshold_kwargs: &threshold_05
         operator: greater
         threshold: 0.5
@@ -53,6 +59,8 @@ metrics:
       metric: precision
       multitask_handling: mean-per-label
       average: micro
+      task: binary
+      target_nan_mask: ignore
       threshold_kwargs: *threshold_05
   zinc: *qm9_metrics
 
diff --git a/expts/hydra-configs/training/toymix.yaml b/expts/hydra-configs/training/toymix.yaml
index 4afcbd56a..dc6c174b9 100644
--- a/expts/hydra-configs/training/toymix.yaml
+++ b/expts/hydra-configs/training/toymix.yaml
@@ -23,4 +23,4 @@ trainer:
     precision: 16
     max_epochs: ${constants.max_epochs}
     min_epochs: 1
-    check_val_every_n_epoch: 20
\ No newline at end of file
+    check_val_every_n_epoch: 2
\ No newline at end of file
diff --git a/graphium/config/_loader.py b/graphium/config/_loader.py
index 2b9055a83..9596e2eac 100644
--- a/graphium/config/_loader.py
+++ b/graphium/config/_loader.py
@@ -15,7 +15,7 @@
 # Misc
 import os
 from copy import deepcopy
-from typing import Any, Callable, Dict, Mapping, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, Mapping, Optional, Tuple, Type, Union, Iterable
 
 import joblib
 import mup
@@ -40,10 +40,10 @@
 from graphium.trainer.metrics import MetricWrapper
 from graphium.trainer.predictor import PredictorModule
 from graphium.utils.command_line_utils import get_anchors_and_aliases, update_config
+from graphium.trainer.progress_bar import ProgressBarMetrics
 
 # Graphium
 from graphium.utils.mup import set_base_shapes
-from graphium.utils.spaces import DATAMODULE_DICT, GRAPHIUM_PRETRAINED_MODELS_DICT
 from graphium.utils import fs
 
 
@@ -111,6 +111,8 @@ def load_datamodule(
         datamodule: The datamodule used to process and load the data
     """
 
+    from graphium.utils.spaces import DATAMODULE_DICT # Avoid circular imports with `spaces.py`
+
     cfg_data = config["datamodule"]["args"]
 
     # Instanciate the datamodule
@@ -298,9 +300,6 @@ def load_predictor(
     accelerator_type: str,
     featurization: Dict[str, str] = None,
     task_norms: Optional[Dict[Callable, Any]] = None,
-    replicas: int = 1,
-    gradient_acc: int = 1,
-    global_bs: int = 1,
 ) -> PredictorModule:
     """
     Defining the predictor module, which handles the training logic from `lightning.LighningModule`
@@ -326,9 +325,6 @@ def load_predictor(
         task_levels=task_levels,
         featurization=featurization,
         task_norms=task_norms,
-        replicas=replicas,
-        gradient_acc=gradient_acc,
-        global_bs=global_bs,
         **cfg_pred,
     )
 
@@ -345,9 +341,6 @@ def load_predictor(
             task_levels=task_levels,
             featurization=featurization,
             task_norms=task_norms,
-            replicas=replicas,
-            gradient_acc=gradient_acc,
-            global_bs=global_bs,
             **cfg_pred,
         )
 
@@ -384,6 +377,7 @@ def load_trainer(
     config: Union[omegaconf.DictConfig, Dict[str, Any]],
     accelerator_type: str,
     date_time_suffix: str = "",
+    metrics_on_progress_bar: Optional[Iterable[str]] = None,
 ) -> Trainer:
     """
     Defining the pytorch-lightning Trainer module.
@@ -449,12 +443,15 @@ def load_trainer(
             name += f"_{date_time_suffix}"
         trainer_kwargs["logger"] = WandbLogger(name=name, **wandb_cfg)
 
-    trainer_kwargs["callbacks"] = callbacks
+    progress_bar_callback = ProgressBarMetrics(metrics_on_progress_bar = metrics_on_progress_bar)
+    callbacks.append(progress_bar_callback)
+
     trainer = Trainer(
         detect_anomaly=True,
         strategy=strategy,
         accelerator=accelerator_type,
         devices=devices,
+        callbacks=callbacks,
         **cfg_trainer["trainer"],
         **trainer_kwargs,
     )
@@ -625,6 +622,8 @@ def get_checkpoint_path(config: Union[omegaconf.DictConfig, Dict[str, Any]]) ->
     Otherwise, assume it refers to a file in the checkpointing dir.
     """
 
+    from graphium.utils.spaces import GRAPHIUM_PRETRAINED_MODELS_DICT # Avoid circular imports with `spaces.py`
+
     cfg_trainer = config["trainer"]
 
     path = config.get("ckpt_name_for_testing", "last.ckpt")
diff --git a/graphium/config/dummy_finetuning_from_gnn.yaml b/graphium/config/dummy_finetuning_from_gnn.yaml
index 4de1e79bc..75848c40f 100644
--- a/graphium/config/dummy_finetuning_from_gnn.yaml
+++ b/graphium/config/dummy_finetuning_from_gnn.yaml
@@ -55,7 +55,7 @@ finetuning:
 
 constants:
   seed: 42
-  max_epochs: 2
+  max_epochs: 5
 
 accelerator:
   float32_matmul_precision: medium
@@ -64,14 +64,14 @@ accelerator:
 predictor:
   random_seed: ${constants.seed}
   optim_kwargs:
-    lr: 4.e-5
+    lr: 1.e-3
   scheduler_kwargs: null
-  target_nan_mask: null
+  target_nan_mask: ignore
   multitask_handling: flatten # flatten, mean-per-label
   
   torch_scheduler_kwargs:
     module_type: WarmUpLinearLR
-    max_num_epochs: 2
+    max_num_epochs: 4
     warmup_epochs: 1
     verbose: False
   
@@ -84,22 +84,22 @@ metrics:
   lipophilicity_astrazeneca:
     - name: mae
       metric: mae
-      target_nan_mask: null
+      target_nan_mask: ignore
       multitask_handling: flatten
       threshold_kwargs: null
     - name: spearman
       metric: spearmanr
       threshold_kwargs: null
-      target_nan_mask: null
+      target_nan_mask: ignore
       multitask_handling: mean-per-label
     - name: pearson
       metric: pearsonr
       threshold_kwargs: null
-      target_nan_mask: null
+      target_nan_mask: ignore
       multitask_handling: mean-per-label
     - name: r2_score
       metric: r2_score
-      target_nan_mask: null
+      target_nan_mask: ignore
       multitask_handling: mean-per-label
       threshold_kwargs: null
 
@@ -107,7 +107,7 @@ trainer:
   seed: ${constants.seed}
   trainer:
     precision: 32
-    max_epochs: 2
+    max_epochs: 5
     min_epochs: 1
     check_val_every_n_epoch: 1
     accumulate_grad_batches: 1
@@ -122,12 +122,12 @@ datamodule:
 
   module_type: "ADMETBenchmarkDataModule"
   args:
+    processed_graph_data_path: datacache/processed_graph_data/dummy_finetuning_from_gnn
     # TDC specific
     tdc_benchmark_names: [lipophilicity_astrazeneca]
     tdc_train_val_seed: ${constants.seed}
-    
-    batch_size_training: 200
-    batch_size_inference: 200
+    batch_size_training: 20
+    batch_size_inference: 20
     num_workers: 0
 
     persistent_workers: False
\ No newline at end of file
diff --git a/graphium/config/dummy_finetuning_from_task_head.yaml b/graphium/config/dummy_finetuning_from_task_head.yaml
index 90b0d5341..373bc6e7e 100644
--- a/graphium/config/dummy_finetuning_from_task_head.yaml
+++ b/graphium/config/dummy_finetuning_from_task_head.yaml
@@ -61,7 +61,7 @@ finetuning:
 
 constants:
   seed: 42
-  max_epochs: 2
+  max_epochs: 5
 
 accelerator:
   float32_matmul_precision: medium
@@ -70,14 +70,14 @@ accelerator:
 predictor:
   random_seed: ${constants.seed}
   optim_kwargs:
-    lr: 4.e-5
+    lr: 1.e-3
   scheduler_kwargs: null
   target_nan_mask: null
   multitask_handling: flatten # flatten, mean-per-label
   
   torch_scheduler_kwargs:
     module_type: WarmUpLinearLR
-    max_num_epochs: 2
+    max_num_epochs: 5
     warmup_epochs: 1
     verbose: False
   
@@ -113,7 +113,7 @@ trainer:
   seed: ${constants.seed}
   trainer:
     precision: 32
-    max_epochs: 2
+    max_epochs: 4
     min_epochs: 1
     check_val_every_n_epoch: 1
     accumulate_grad_batches: 1
@@ -128,12 +128,12 @@ datamodule:
 
   module_type: "ADMETBenchmarkDataModule"
   args:
+    processed_graph_data_path: datacache/processed_graph_data/dummy_finetuning_task_head
     # TDC specific
     tdc_benchmark_names: [lipophilicity_astrazeneca]
     tdc_train_val_seed: ${constants.seed}
-    
-    batch_size_training: 200
-    batch_size_inference: 200
+    batch_size_training: 20
+    batch_size_inference: 20
     num_workers: 0
 
     persistent_workers: False
diff --git a/graphium/data/datamodule.py b/graphium/data/datamodule.py
index e2368d3b3..edc5e207d 100644
--- a/graphium/data/datamodule.py
+++ b/graphium/data/datamodule.py
@@ -845,6 +845,7 @@ def __init__(
         )
         IPUDataModuleModifier.__init__(self, **kwargs)
 
+        self._len = None
         self.task_specific_args = task_specific_args
 
         self.task_dataset_processing_params = {}
@@ -929,6 +930,30 @@ def encode_feature_options(options, name, encoding_function):
 
         if self._ready_to_load_all_from_file():
             self._data_is_prepared = True
+            self._len = self._get_len_from_cached_file()
+
+    def _get_len_from_cached_file(self):
+        if self._ready_to_load_all_from_file():
+            self._data_is_prepared = True
+            train_metadata = graphium_cpp.load_metadata_tensors(
+                        self.processed_graph_data_path, "train", self.data_hash
+                    )
+            val_metadata = graphium_cpp.load_metadata_tensors(
+                    self.processed_graph_data_path, "val", self.data_hash
+                )
+            test_metadata = graphium_cpp.load_metadata_tensors(
+                    self.processed_graph_data_path, "test", self.data_hash
+                )
+            length = 0
+            if len(train_metadata) > 0:
+                length += len(train_metadata[2])
+            if len(val_metadata) > 0:
+                length += len(val_metadata[2])
+            if len(test_metadata) > 0:
+                length += len(test_metadata[2])
+        else:
+            raise ValueError("Data is not prepared. Please call prepare_data() first.")
+        return length
 
     def _parse_caching_args(self, processed_graph_data_path):
         """
@@ -1102,6 +1127,7 @@ def prepare_data(self):
                 num_molecules,
                 split_val=self.task_dataset_processing_params[task].split_val,
                 split_test=self.task_dataset_processing_params[task].split_test,
+                sample_idx=sample_idx,
                 split_seed=self.task_dataset_processing_params[task].seed,
                 splits_path=self.task_dataset_processing_params[task].splits_path,
                 split_names=self.task_dataset_processing_params[task].split_names,
@@ -1138,6 +1164,7 @@ def prepare_data(self):
             self.explicit_H,
             self.preprocessing_n_jobs,
         )
+        self._len = self._get_len_from_cached_file()
 
         for task, stats in all_stats.items():
             if len(stats) < 4:
@@ -1718,18 +1745,14 @@ def get_data_hash(self):
 
     def __len__(self) -> int:
         r"""
-        Returns the number of elements of the current DataModule, which is the combined size of all single-task datasets given.
+        Returns the number of smiles of the current DataModule, which depends on all the smiles from all tasks.
+        If `prepare_data` is not called, the length is unknown and will raise an error.
         Returns:
             num_elements: Number of elements in the current DataModule
         """
-        num_elements = 0
-        for task, args in self.task_dataset_processing_params.items():
-            if args.df is None:
-                df = self._read_table(args.df_path, usecols=[args.smiles_col])
-                num_elements += len(df)
-            else:
-                num_elements += len(args.df)
-        return num_elements
+        if self._len is None:
+            raise ValueError("The length of the dataset is unknown. Please call `prepare_data` first.")
+        return self._len
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -2075,6 +2098,12 @@ def __init__(
             for t in tdc_benchmark_names
         }
 
+        # Create a temporary `processed_graph_data_path` to store the processed graphs and labels
+        if processed_graph_data_path is None:
+            processed_graph_data_path = fs.join(tdc_cache_dir, "processed_graph_data")
+            if not fs.exists(processed_graph_data_path):
+                fs.mkdir(processed_graph_data_path)
+
         super().__init__(
             task_specific_args=task_specific_args,
             featurization=featurization,
diff --git a/graphium/finetuning/finetuning.py b/graphium/finetuning/finetuning.py
index 97d6d7fc7..59902a921 100644
--- a/graphium/finetuning/finetuning.py
+++ b/graphium/finetuning/finetuning.py
@@ -17,10 +17,10 @@
 from collections import OrderedDict
 
 import torch.nn as nn
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 
 from torch.optim.optimizer import Optimizer
-from pytorch_lightning.callbacks import BaseFinetuning
+from lightning.pytorch.callbacks import BaseFinetuning
 
 
 class GraphFinetuning(BaseFinetuning):
diff --git a/graphium/finetuning/finetuning_architecture.py b/graphium/finetuning/finetuning_architecture.py
index 864016141..4b0de1607 100644
--- a/graphium/finetuning/finetuning_architecture.py
+++ b/graphium/finetuning/finetuning_architecture.py
@@ -21,7 +21,6 @@
 
 from graphium.nn.utils import MupMixin
 from graphium.trainer.predictor import PredictorModule
-from graphium.utils.spaces import FINETUNING_HEADS_DICT
 
 
 class FullGraphFinetuningNetwork(nn.Module, MupMixin):
@@ -309,6 +308,8 @@ def __init__(self, finetuning_head_kwargs: Dict[str, Any]):
 
         """
 
+        from graphium.utils.spaces import FINETUNING_HEADS_DICT # Avoiding circular imports with `spaces.py`
+        
         super().__init__()
         self.task = finetuning_head_kwargs.pop("task", None)
         self.previous_module = finetuning_head_kwargs.pop("previous_module", "task_heads")
diff --git a/graphium/graphium_cpp/pybind11 b/graphium/graphium_cpp/pybind11
new file mode 160000
index 000000000..ccefee4c3
--- /dev/null
+++ b/graphium/graphium_cpp/pybind11
@@ -0,0 +1,1 @@
+Subproject commit ccefee4c3187c2892fcf4590b1bbc850134b84bb
diff --git a/graphium/ipu/ipu_metrics.py b/graphium/ipu/ipu_metrics.py
deleted file mode 100644
index 9029d3e00..000000000
--- a/graphium/ipu/ipu_metrics.py
+++ /dev/null
@@ -1,907 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-from typing import Optional, Tuple, Sequence, Literal
-
-import torch
-from torch import BoolTensor, IntTensor, Tensor
-from torchmetrics.functional import auroc, average_precision, pearson_corrcoef, r2_score
-from torchmetrics.utilities.checks import _input_squeeze
-from torchmetrics.functional.classification.accuracy import (
-    _mode,
-    _check_subset_validity,
-    _accuracy_compute,
-    _accuracy_update,
-)
-from torchmetrics.functional.classification.precision_recall import _precision_compute, _recall_compute
-from torchmetrics.functional.classification.f_beta import _fbeta_compute
-from torchmetrics.functional import mean_squared_error, mean_absolute_error
-from torchmetrics.utilities.checks import _input_squeeze
-from torchmetrics.utilities.enums import AverageMethod
-
-from graphium.utils.tensor import nan_mean
-from graphium.ipu.ipu_utils import import_poptorch
-
-
-def auroc_ipu(
-    preds: Tensor,
-    target: Tensor,
-    num_classes: Optional[int] = None,
-    task: Optional[Literal["binary", "multiclass", "multilabel"]] = None,
-    pos_label: Optional[int] = None,
-    average: Optional[str] = "macro",
-    max_fpr: Optional[float] = None,
-    sample_weights: Optional[Sequence] = None,
-):
-    """
-    A modified version of the `torchmetrics.functional.auroc` that can ignore NaNs
-    by giving them the same value for both `preds` and `target`.
-    This allows it to work with compilation
-    and IPUs since it doesn't modify the tensor's shape.
-    """
-
-    target = target.clone()
-    preds = preds.clone()
-
-    # Replace the nan-targets in the preds/target tensors by 0
-    nan_targets = target.isnan()
-    preds[nan_targets] = 0.0
-    target[nan_targets] = 0.0
-
-    # Get the original weight matrix. If None, set all weights = 1
-    if sample_weights is None:
-        sample_weights = torch.ones(target.shape[0], dtype=preds.dtype, device=preds.device)
-    sample_weights[nan_targets] = 0.0
-
-    # Compute the loss, and rescale by the number of nan elements
-    score = auroc(
-        preds=preds,
-        target=target.to(int),
-        num_classes=num_classes,
-        task=task,
-        pos_label=pos_label,
-        average=average,
-        max_fpr=max_fpr,
-        sample_weights=sample_weights,
-    )
-
-    return score
-
-
-def average_precision_ipu(
-    preds: Tensor,
-    target: Tensor,
-    num_classes: Optional[int] = None,
-    task: Optional[Literal["binary", "multiclass", "multilabel"]] = None,
-    ignore_index: Optional[int] = None,
-    pos_label: Optional[int] = None,
-    average: Optional[str] = "macro",
-    sample_weights: Optional[Sequence] = None,
-):
-    """
-    A modified version of the `torchmetrics.functional.average_precision` that can ignore NaNs
-    by giving them the same value for both `preds` and `target`.
-    This allows it to work with compilation
-    and IPUs since it doesn't modify the tensor's shape.
-    """
-
-    target = target.clone()
-    preds = preds.clone()
-
-    # Replace the nan-targets in the preds/target tensors by 0
-    # Average precision is not sensitive to true negatives
-    nan_targets = target.isnan()
-    preds[nan_targets] = 0.0
-    target[nan_targets] = 0.0
-
-    # No need to use sample weights (which is no longer supported in torchmetrics >=0.10)
-    # # Get the original weight matrix. If None, set all weights = 1
-    # if sample_weights is None:
-    #     sample_weights = torch.ones(target.shape[0], dtype=preds.dtype, device=preds.device)
-    # sample_weights[nan_targets] = 0.0
-
-    # Compute the loss, and rescale by the number of nan elements
-    score = average_precision(
-        preds=preds,
-        target=target,
-        num_classes=num_classes,
-        task=task,
-        ignore_index=ignore_index,
-        pos_label=pos_label,
-        average=average,
-        # sample_weights=sample_weights,
-    )
-
-    return score
-
-
-def precision_ipu(
-    preds: Tensor,
-    target: Tensor,
-    average: Optional[str] = "micro",
-    mdmc_average: Optional[str] = None,
-    ignore_index: Optional[int] = None,
-    num_classes: Optional[int] = None,
-    threshold: float = 0.5,
-    top_k: Optional[int] = None,
-    multiclass: Optional[bool] = None,
-):
-    """
-    A modified version of the `torchmetrics.functional.precision` that can ignore NaNs
-    by giving them the same value for both `preds` and `target`.
-    This allows it to work with compilation
-    and IPUs since it doesn't modify the tensor's shape.
-    """
-
-    (tp, fp, tn, fn), mode = get_confusion_matrix(
-        preds=preds,
-        target=target,
-        average=average,
-        mdmc_average=mdmc_average,
-        threshold=threshold,
-        top_k=top_k,
-        subset_accuracy=False,
-        num_classes=num_classes,
-        multiclass=multiclass,
-        ignore_index=ignore_index,
-    )
-
-    return _precision_compute(tp, fp, fn, average, mdmc_average)
-
-
-def recall_ipu(
-    preds: Tensor,
-    target: Tensor,
-    average: Optional[str] = "micro",
-    mdmc_average: Optional[str] = None,
-    ignore_index: Optional[int] = None,
-    num_classes: Optional[int] = None,
-    threshold: float = 0.5,
-    top_k: Optional[int] = None,
-    multiclass: Optional[bool] = None,
-):
-    """
-    A modified version of the `torchmetrics.functional.recall` that can ignore NaNs
-    by giving them the same value for both `preds` and `target`.
-    This allows it to work with compilation
-    and IPUs since it doesn't modify the tensor's shape.
-    """
-
-    (tp, fp, tn, fn), mode = get_confusion_matrix(
-        preds=preds,
-        target=target,
-        average=average,
-        mdmc_average=mdmc_average,
-        threshold=threshold,
-        top_k=top_k,
-        num_classes=num_classes,
-        multiclass=multiclass,
-        ignore_index=ignore_index,
-    )
-
-    return _recall_compute(tp, fp, fn, average, mdmc_average)
-
-
-def accuracy_ipu(
-    preds: Tensor,
-    target: Tensor,
-    average: Optional[str] = "micro",
-    mdmc_average: Optional[str] = "global",
-    threshold: float = 0.5,
-    top_k: Optional[int] = None,
-    subset_accuracy: bool = False,
-    num_classes: Optional[int] = None,
-    multiclass: Optional[bool] = None,
-    ignore_index: Optional[int] = None,
-) -> Tensor:
-    """
-    A modified version of the `torchmetrics.functional.accuracy` that can ignore NaNs
-    by giving them the same value for both `preds` and `target`.
-    This allows it to work with compilation
-    and IPUs since it doesn't modify the tensor's shape.
-
-    Args:
-        preds: Predictions from model (probabilities, logits or labels)
-        target: Ground truth labels
-        average:
-            Defines the reduction that is applied. Should be one of the following:
-
-            - ``'micro'`` [default]: Calculate the metric globally, across all samples and classes.
-            - ``'macro'``: Calculate the metric for each class separately, and average the
-              metrics across classes (with equal weights for each class).
-            - ``'weighted'``: Calculate the metric for each class separately, and average the
-              metrics across classes, weighting each class by its support (``tp + fn``).
-            - ``'none'`` or ``None``: Calculate the metric for each class separately, and return
-              the metric for every class.
-            - ``'samples'``: Calculate the metric for each sample, and average the metrics
-              across samples (with equal weights for each sample).
-
-            .. note:: What is considered a sample in the multi-dimensional multi-class case
-                depends on the value of ``mdmc_average``.
-
-            .. note:: If ``'none'`` and a given class doesn't occur in the ``preds`` or ``target``,
-                the value for the class will be ``nan``.
-
-        mdmc_average:
-            Defines how averaging is done for multi-dimensional multi-class inputs (on top of the
-            ``average`` parameter). Should be one of the following:
-
-            - ``None`` [default]: Should be left unchanged if your data is not multi-dimensional multi-class.
-
-            - ``'samplewise'``: In this case, the statistics are computed separately for each
-              sample on the ``N`` axis, and then averaged over samples.
-              The computation for each sample is done by treating the flattened extra axes ``...``
-              (see :ref:`pages/classification:input types`) as the ``N`` dimension within the sample,
-              and computing the metric for the sample based on that.
-
-            - ``'global'``: In this case the ``N`` and ``...`` dimensions of the inputs
-              (see :ref:`pages/classification:input types`)
-              are flattened into a new ``N_X`` sample axis, i.e. the inputs are treated as if they
-              were ``(N_X, C)``. From here on the ``average`` parameter applies as usual.
-
-        num_classes:
-            Number of classes. Necessary for ``'macro'``, ``'weighted'`` and ``None`` average methods.
-
-        threshold:
-            Threshold for transforming probability or logit predictions to binary (0,1) predictions, in the case
-            of binary or multi-label inputs. Default value of 0.5 corresponds to input being probabilities.
-        top_k:
-            Number of the highest probability or logit score predictions considered finding the correct label,
-            relevant only for (multi-dimensional) multi-class inputs. The
-            default value (``None``) will be interpreted as 1 for these inputs.
-
-            Should be left at default (``None``) for all other types of inputs.
-        multiclass:
-            Used only in certain special cases, where you want to treat inputs as a different type
-            than what they appear to be. See the parameter's
-            :ref:`documentation section <pages/classification:using the multiclass parameter>`
-            for a more detailed explanation and examples.
-        ignore_index:
-            Integer specifying a target class to ignore. If given, this class index does not contribute
-            to the returned score, regardless of reduction method. If an index is ignored, and ``average=None``
-            or ``'none'``, the score for the ignored class will be returned as ``nan``.
-        subset_accuracy:
-            Whether to compute subset accuracy for multi-label and multi-dimensional
-            multi-class inputs (has no effect for other input types).
-
-            - For multi-label inputs, if the parameter is set to ``True``, then all labels for
-              each sample must be correctly predicted for the sample to count as correct. If it
-              is set to ``False``, then all labels are counted separately - this is equivalent to
-              flattening inputs beforehand (i.e. ``preds = preds.flatten()`` and same for ``target``).
-
-            - For multi-dimensional multi-class inputs, if the parameter is set to ``True``, then all
-              sub-sample (on the extra axis) must be correct for the sample to be counted as correct.
-              If it is set to ``False``, then all sub-samples are counter separately - this is equivalent,
-              in the case of label predictions, to flattening the inputs beforehand (i.e.
-              ``preds = preds.flatten()`` and same for ``target``). Note that the ``top_k`` parameter
-              still applies in both cases, if set.
-
-    Raises:
-        ValueError:
-            If ``top_k`` parameter is set for ``multi-label`` inputs.
-        ValueError:
-            If ``average`` is none of ``"micro"``, ``"macro"``, ``"weighted"``, ``"samples"``, ``"none"``, ``None``.
-        ValueError:
-            If ``mdmc_average`` is not one of ``None``, ``"samplewise"``, ``"global"``.
-        ValueError:
-            If ``average`` is set but ``num_classes`` is not provided.
-        ValueError:
-            If ``num_classes`` is set
-            and ``ignore_index`` is not in the range ``[0, num_classes)``.
-        ValueError:
-            If ``top_k`` is not an ``integer`` larger than ``0``.
-    """
-
-    (tp, fp, tn, fn), mode = get_confusion_matrix(
-        preds=preds,
-        target=target,
-        average=average,
-        mdmc_average=mdmc_average,
-        threshold=threshold,
-        top_k=top_k,
-        subset_accuracy=subset_accuracy,
-        num_classes=num_classes,
-        multiclass=multiclass,
-        ignore_index=ignore_index,
-    )
-
-    return _accuracy_compute(tp, fp, tn, fn, average, mdmc_average, mode)
-
-
-def get_confusion_matrix(
-    preds: Tensor,
-    target: Tensor,
-    average: Optional[str] = "micro",
-    mdmc_average: Optional[str] = "global",
-    threshold: float = 0.5,
-    top_k: Optional[int] = None,
-    subset_accuracy: bool = False,
-    num_classes: Optional[int] = None,
-    multiclass: Optional[bool] = None,
-    ignore_index: Optional[int] = None,
-) -> Tuple[Tuple[Tensor], Tensor]:
-    """
-    Calculates the confusion matrix according to the specified average method.
-
-    Args:
-        preds: Predictions from model (probabilities, logits or labels)
-        target: Ground truth labels
-        average:
-            Defines the reduction that is applied. Should be one of the following:
-
-            - ``'micro'`` [default]: Calculate the metric globally, across all samples and classes.
-            - ``'macro'``: Calculate the metric for each class separately, and average the
-              metrics across classes (with equal weights for each class).
-            - ``'weighted'``: Calculate the metric for each class separately, and average the
-              metrics across classes, weighting each class by its support (``tp + fn``).
-            - ``'none'`` or ``None``: Calculate the metric for each class separately, and return
-              the metric for every class.
-            - ``'samples'``: Calculate the metric for each sample, and average the metrics
-              across samples (with equal weights for each sample).
-
-            .. note:: What is considered a sample in the multi-dimensional multi-class case
-                depends on the value of ``mdmc_average``.
-
-            .. note:: If ``'none'`` and a given class doesn't occur in the ``preds`` or ``target``,
-                the value for the class will be ``nan``.
-
-        mdmc_average:
-            Defines how averaging is done for multi-dimensional multi-class inputs (on top of the
-            ``average`` parameter). Should be one of the following:
-
-            - ``None`` [default]: Should be left unchanged if your data is not multi-dimensional multi-class.
-
-            - ``'samplewise'``: In this case, the statistics are computed separately for each
-              sample on the ``N`` axis, and then averaged over samples.
-              The computation for each sample is done by treating the flattened extra axes ``...``
-              (see :ref:`pages/classification:input types`) as the ``N`` dimension within the sample,
-              and computing the metric for the sample based on that.
-
-            - ``'global'``: In this case the ``N`` and ``...`` dimensions of the inputs
-              (see :ref:`pages/classification:input types`)
-              are flattened into a new ``N_X`` sample axis, i.e. the inputs are treated as if they
-              were ``(N_X, C)``. From here on the ``average`` parameter applies as usual.
-
-        num_classes:
-            Number of classes. Necessary for ``'macro'``, ``'weighted'`` and ``None`` average methods.
-
-        threshold:
-            Threshold for transforming probability or logit predictions to binary (0,1) predictions, in the case
-            of binary or multi-label inputs. Default value of 0.5 corresponds to input being probabilities.
-        top_k:
-            Number of the highest probability or logit score predictions considered finding the correct label,
-            relevant only for (multi-dimensional) multi-class inputs. The
-            default value (``None``) will be interpreted as 1 for these inputs.
-
-            Should be left at default (``None``) for all other types of inputs.
-        multiclass:
-            Used only in certain special cases, where you want to treat inputs as a different type
-            than what they appear to be. See the parameter's
-            :ref:`documentation section <pages/classification:using the multiclass parameter>`
-            for a more detailed explanation and examples.
-        ignore_index:
-            Integer specifying a target class to ignore. If given, this class index does not contribute
-            to the returned score, regardless of reduction method. If an index is ignored, and ``average=None``
-    """
-    allowed_average = ["micro", "macro", "weighted", "samples", "none", None]
-    if average not in allowed_average:
-        raise ValueError(f"The `average` has to be one of {allowed_average}, got {average}.")
-
-    if average in ["macro", "weighted", "none", None] and (not num_classes or num_classes < 1):
-        raise ValueError(f"When you set `average` as {average}, you have to provide the number of classes.")
-
-    allowed_mdmc_average = [None, "samplewise", "global"]
-    if mdmc_average not in allowed_mdmc_average:
-        raise ValueError(f"The `mdmc_average` has to be one of {allowed_mdmc_average}, got {mdmc_average}.")
-
-    if num_classes and ignore_index is not None and (not ignore_index < num_classes or num_classes == 1):
-        raise ValueError(
-            f"The `ignore_index` {ignore_index} is not valid for inputs with {num_classes} classes"
-        )
-
-    if top_k is not None and (not isinstance(top_k, int) or top_k <= 0):
-        raise ValueError(f"The `top_k` should be an integer larger than 0, got {top_k}")
-
-    #### ADDED ####
-    # Put all the NaNs as the 0-class
-    nans = torch.isnan(target)
-    target[nans] = 0
-    preds[nans] = 0
-    if (preds.ndim > 1) and (preds.shape[1] > 1):
-        preds[nans, 0] = 1
-    target = target.to(int)
-    #### END ADDED ####
-
-    preds, target = _input_squeeze(preds, target)
-    mode = _mode(preds, target, threshold, top_k, num_classes, multiclass, ignore_index)
-    reduce = "macro" if average in ["weighted", "none", None] else average
-
-    if subset_accuracy and _check_subset_validity(mode):
-        # correct, total = _subset_accuracy_update(preds, target, threshold, top_k, ignore_index)
-        # return _subset_accuracy_compute(correct, total)
-        raise NotImplementedError("subset_accuracy not implemented")
-    tp, fp, tn, fn = _accuracy_update(
-        preds, target, reduce, mdmc_average, threshold, num_classes, top_k, multiclass, ignore_index, mode
-    )
-
-    #### ADDED ####
-    num_nans = nans.sum(0)
-    if tp.numel() > 1:
-        tp[0] = tp[0] - num_nans
-        tn[1:] = tn[1:] - num_nans
-    else:
-        tn = tn - num_nans
-        if (preds.ndim > 1) and (preds.shape[1] > 1):
-            tp = tp - num_nans
-    #### END ADDED ####
-
-    return (tp, fp, tn, fn), mode
-
-
-class NaNTensor(Tensor):
-    """
-    Class to create and manage a NaN tensor along it's properties
-
-    The goal of the class is to override the regular tensor such that the basic
-    operations (sum, mean, max, etc) ignore the NaNs in the input.
-    It also supports NaNs in integer tensors (as the lowest integer possible).
-    """
-
-    @property
-    def get_nans(self) -> BoolTensor:
-        """
-        Gets the boolean Tensor containing the location of NaNs.
-        In the case of an integer tensor, this returns where the tensor is equal to its minimal value
-        In the case of a boolean tensor, this returns a Tensor filled with `False`
-        """
-        if self.is_floating_point():
-            return self.isnan()
-        elif self.is_signed():
-            return self == torch.iinfo(self.dtype).min
-        else:
-            return torch.zeros(self.shape, device=self.device, dtype=bool)
-
-    def sum(self, *args, **kwargs) -> Tensor:
-        """
-        Overloads the traditional sum to ignore the NaNs
-        """
-        tensor = self.to(float)
-        tensor[self.get_nans] = float("nan")
-        if self.is_floating_point():
-            dtype = self.dtype
-        else:
-            dtype = torch.int64
-        return tensor.nansum(*args, **kwargs).to(dtype)
-
-    def mean(self, *args, **kwargs) -> Tensor:
-        """
-        Overloads the traditional mean to ignore the NaNs
-        """
-        tensor = self.to(float)
-        tensor[self.get_nans] = float("nan")
-        return nan_mean(tensor, *args, **kwargs).to(self.dtype)
-
-    def numel(self) -> int:
-        """
-        Returns the number of non-NaN elements.
-        """
-        return super(NaNTensor, ~self.get_nans).sum()
-
-    def min(self, *args, **kwargs) -> Tensor:
-        """
-        Returns the min vale of a tensor whitout NaNs
-        """
-        tensor = self
-        tensor = tensor[~self.get_nans]
-        return super(NaNTensor, tensor).min(*args, **kwargs)
-
-    def max(self, *args, **kwargs) -> Tensor:
-        """
-        Returns the max vale of a tensor whitout NaNs
-        """
-        tensor = self
-        tensor = tensor[~self.get_nans]
-        return super(NaNTensor, tensor).max(*args, **kwargs)
-
-    def argsort(self, dim=-1, descending=False) -> IntTensor:
-        """
-        Return the indices that sort the tensor, while putting all the NaNs to the end of the sorting.
-        """
-        tensor = self
-        if descending:
-            tensor[tensor.get_nans] = float("-inf")
-        else:
-            tensor[tensor.get_nans] = float("inf")
-        return super(NaNTensor, tensor).argsort(dim=dim, descending=descending)
-
-    def size(self, dim) -> Tensor:
-        """
-        Instead of returning the size, return the number of non-NaN elements in
-        a specific dimension. Useful for the `r2_score` metric.
-        """
-        return (~self.get_nans).sum(dim=dim)
-
-    def __lt__(self, other) -> Tensor:
-        """
-        Stupid fix that allows the code to work with `r2_score`,
-        since it requires the size to be > 2. But since `self.size` now returns
-        a Tensor instead of a value, we check that all elements are > 2.
-        """
-        if (not isinstance(other, Tensor)) and (other == 2):
-            return super().__lt__(other).all()
-        else:
-            return super().__lt__(other)
-
-    @classmethod
-    def __torch_function__(cls, func, types, args=(), kwargs=None):
-        """
-        This __torch_function__ implementation wraps subclasses such that
-        methods called on subclasses return a subclass instance instead of
-        a ``torch.Tensor`` instance.
-
-        One corollary to this is that you need coverage for torch.Tensor
-        methods if implementing __torch_function__ for subclasses.
-
-        Affects the call torch.sum() as to behave the same way as NaNTensor.sum()
-
-        We recommend always calling ``super().__torch_function__`` as the base
-        case when doing the above.
-
-        While not mandatory, we recommend making `__torch_function__` a classmethod.
-        """
-        if func.__name__ == "sum":
-            kwargs = {} if kwargs is None else kwargs
-            return args[0].sum(*args[1:], **kwargs)
-        else:
-            return super().__torch_function__(func, types, args=args, kwargs=kwargs)
-
-
-def pearson_ipu(preds, target):
-    """Computes pearson correlation coefficient.
-
-    Handles NaNs in the target without reshaping tensors in order to work on IPU.
-
-    Args:
-        preds: estimated scores
-        target: ground truth scores
-    """
-    preds = NaNTensor(preds)
-    target = NaNTensor(target)
-    preds[target.get_nans] = float("nan")
-    pearson = pearson_corrcoef(preds, target.to(preds.dtype))
-    return Tensor(pearson)
-
-
-def spearman_ipu(preds, target):
-    """Computes spearman rank correlation coefficient.
-
-    Handles NaNs in the target without reshaping tensors in order to work on IPU.
-
-    Args:
-        preds: estimated scores
-        target: ground truth scores
-    """
-    nans = target.isnan()
-    dtype = preds.dtype
-    preds[nans] = float("inf")
-    target[nans] = float("inf")
-    preds_sort = _rank_data(preds).to(dtype=dtype)
-    target_sort = _rank_data(target).to(dtype=dtype)
-    target_sort[nans] = float("nan")
-    spearman = pearson_ipu(preds_sort, target_sort)
-    return Tensor(spearman)
-
-
-def _rank_data(data: Tensor) -> Tensor:
-    """Calculate the rank for each element of a tensor.
-
-    The rank refers to the indices of an element in the corresponding sorted tensor (starting from 1).
-    Duplicates of the same value will be assigned the mean of their rank.
-
-    Adopted from `Rank of element tensor`_
-    """
-    n = data.numel()
-    rank = torch.empty_like(data)
-    idx = data.argsort()
-    rank[idx] = torch.arange(1, n + 1, dtype=data.dtype, device=data.device)
-
-    # TODO: Repeats not yet supported
-    # repeats = _find_repeats(data)
-    # for r in repeats:
-    #     condition = data == r
-    #     rank[condition] = rank[condition].mean()
-    return rank
-
-
-def r2_score_ipu(preds, target, *args, **kwargs) -> Tensor:
-    """
-    Computes r2 score also known as `R2 Score_Coefficient Determination`_:
-
-    .. math:: R^2 = 1 - \frac{SS_{res}}{SS_{tot}}
-
-    where :math:`SS_{res}=\sum_i (y_i - f(x_i))^2` is the sum of residual squares, and
-    :math:`SS_{tot}=\sum_i (y_i - \bar{y})^2` is total sum of squares. Can also calculate
-    adjusted r2 score given by
-
-    .. math:: R^2_{adj} = 1 - \frac{(1-R^2)(n-1)}{n-k-1}
-
-    where the parameter :math:`k` (the number of independent regressors) should
-    be provided as the ``adjusted`` argument.
-    Handles NaNs without reshaping tensors in order to work on IPU.
-
-    Args:
-        preds: estimated labels
-        target: ground truth labels
-        adjusted: number of independent regressors for calculating adjusted r2 score.
-        multioutput: Defines aggregation in the case of multiple output scores. Can be one of the following strings:
-
-            * ``'raw_values'`` returns full set of scores
-            * ``'uniform_average'`` scores are uniformly averaged
-            * ``'variance_weighted'`` scores are weighted by their individual variances
-    """
-    preds = NaNTensor(preds)
-    target = NaNTensor(target)
-    preds[target.get_nans] = float("nan")
-    score = r2_score(preds, target, *args, **kwargs)
-    return Tensor(score)
-
-
-def fbeta_score_ipu(
-    preds: Tensor,
-    target: Tensor,
-    beta: float = 1.0,
-    average: Optional[str] = "micro",
-    mdmc_average: Optional[str] = None,
-    ignore_index: Optional[int] = None,
-    num_classes: Optional[int] = None,
-    threshold: float = 0.5,
-    top_k: Optional[int] = None,
-    multiclass: Optional[bool] = None,
-):
-    """
-    A modified version of the `torchmetrics.functional.classification.f_beta._fbeta_compute`
-    that can ignore NaNs by giving them the same value for both `preds` and `target`.
-    This allows it to work with compilation
-    and IPUs since it doesn't modify the tensor's shape.
-
-    Args:
-        preds: Predictions from model (probabilities, logits or labels)
-        target: Ground truth labels
-        average:
-            Defines the reduction that is applied. Should be one of the following:
-
-            - ``'micro'`` [default]: Calculate the metric globally, across all samples and classes.
-            - ``'macro'``: Calculate the metric for each class separately, and average the
-              metrics across classes (with equal weights for each class).
-            - ``'weighted'``: Calculate the metric for each class separately, and average the
-              metrics across classes, weighting each class by its support (``tp + fn``).
-            - ``'none'`` or ``None``: Calculate the metric for each class separately, and return
-              the metric for every class.
-            - ``'samples'``: Calculate the metric for each sample, and average the metrics
-              across samples (with equal weights for each sample).
-
-            .. note:: What is considered a sample in the multi-dimensional multi-class case
-                depends on the value of ``mdmc_average``.
-
-            .. note:: If ``'none'`` and a given class doesn't occur in the ``preds`` or ``target``,
-                the value for the class will be ``nan``.
-
-        mdmc_average:
-            Defines how averaging is done for multi-dimensional multi-class inputs (on top of the
-            ``average`` parameter). Should be one of the following:
-
-            - ``None`` [default]: Should be left unchanged if your data is not multi-dimensional multi-class.
-
-            - ``'samplewise'``: In this case, the statistics are computed separately for each
-              sample on the ``N`` axis, and then averaged over samples.
-              The computation for each sample is done by treating the flattened extra axes ``...``
-              (see :ref:`pages/classification:input types`) as the ``N`` dimension within the sample,
-              and computing the metric for the sample based on that.
-
-            - ``'global'``: In this case the ``N`` and ``...`` dimensions of the inputs
-              (see :ref:`pages/classification:input types`)
-              are flattened into a new ``N_X`` sample axis, i.e. the inputs are treated as if they
-              were ``(N_X, C)``. From here on the ``average`` parameter applies as usual.
-
-        num_classes:
-            Number of classes. Necessary for ``'macro'``, ``'weighted'`` and ``None`` average methods.
-
-        threshold:
-            Threshold for transforming probability or logit predictions to binary (0,1) predictions, in the case
-            of binary or multi-label inputs. Default value of 0.5 corresponds to input being probabilities.
-        top_k:
-            Number of the highest probability or logit score predictions considered finding the correct label,
-            relevant only for (multi-dimensional) multi-class inputs. The
-            default value (``None``) will be interpreted as 1 for these inputs.
-
-            Should be left at default (``None``) for all other types of inputs.
-        multiclass:
-            Used only in certain special cases, where you want to treat inputs as a different type
-            than what they appear to be. See the parameter's
-            :ref:`documentation section <pages/classification:using the multiclass parameter>`
-            for a more detailed explanation and examples.
-        ignore_index:
-            Integer specifying a target class to ignore. If given, this class index does not contribute
-            to the returned score, regardless of reduction method. If an index is ignored, and ``average=None``
-            or ``'none'``, the score for the ignored class will be returned as ``nan``.
-        subset_accuracy:
-            Whether to compute subset accuracy for multi-label and multi-dimensional
-            multi-class inputs (has no effect for other input types).
-
-            - For multi-label inputs, if the parameter is set to ``True``, then all labels for
-              each sample must be correctly predicted for the sample to count as correct. If it
-              is set to ``False``, then all labels are counted separately - this is equivalent to
-              flattening inputs beforehand (i.e. ``preds = preds.flatten()`` and same for ``target``).
-
-            - For multi-dimensional multi-class inputs, if the parameter is set to ``True``, then all
-              sub-sample (on the extra axis) must be correct for the sample to be counted as correct.
-              If it is set to ``False``, then all sub-samples are counter separately - this is equivalent,
-              in the case of label predictions, to flattening the inputs beforehand (i.e.
-              ``preds = preds.flatten()`` and same for ``target``). Note that the ``top_k`` parameter
-              still applies in both cases, if set.
-
-    Raises:
-        ValueError:
-            If ``top_k`` parameter is set for ``multi-label`` inputs.
-        ValueError:
-            If ``average`` is none of ``"micro"``, ``"macro"``, ``"weighted"``, ``"samples"``, ``"none"``, ``None``.
-        ValueError:
-            If ``mdmc_average`` is not one of ``None``, ``"samplewise"``, ``"global"``.
-        ValueError:
-            If ``average`` is set but ``num_classes`` is not provided.
-        ValueError:
-            If ``num_classes`` is set
-            and ``ignore_index`` is not in the range ``[0, num_classes)``.
-        ValueError:
-            If ``top_k`` is not an ``integer`` larger than ``0``.
-    """
-
-    (tp, fp, tn, fn), mode = get_confusion_matrix(
-        preds=preds,
-        target=target,
-        average=average,
-        mdmc_average=mdmc_average,
-        ignore_index=ignore_index,
-        num_classes=num_classes,
-        threshold=threshold,
-        top_k=top_k,
-        multiclass=multiclass,
-    )
-
-    b2 = beta**2
-    fbeta = ((1 + b2) * tp) / ((1 + b2) * tp + b2 * fn + fp)
-
-    if average in (None, "none", AverageMethod.NONE):
-        pass
-    elif average == AverageMethod.MICRO:
-        pass
-    elif average == AverageMethod.MACRO:
-        fbeta = fbeta.mean()
-    elif average == AverageMethod.WEIGHTED:
-        weights = tp + fn
-        fbeta = (weights * fbeta).sum() / weights.sum()
-    else:
-        raise ValueError(
-            f"`average={average}` not yet supported. Chose between None, Micro, Macro, or Weighted"
-        )
-
-    return fbeta
-
-
-def f1_score_ipu(
-    preds: Tensor,
-    target: Tensor,
-    beta: float = 1.0,
-    average: Optional[str] = "micro",
-    mdmc_average: Optional[str] = None,
-    ignore_index: Optional[int] = None,
-    num_classes: Optional[int] = None,
-    threshold: float = 0.5,
-    top_k: Optional[int] = None,
-    multiclass: Optional[bool] = None,
-):
-    """
-    A modified version of the `torchmetrics.functional.classification.f_beta._fbeta_compute`
-    that can ignore NaNs by giving them the same value for both `preds` and `target`.
-    Used to calculate the f1_score on IPU with beta parameter equal to 1.0
-    This allows it to work with compilation and IPUs since it doesn't modify the tensor's shape.
-
-    Computes f_beta metric from stat scores: true positives, false positives, true negatives, false negatives.
-
-    Args:
-        tp: True positives
-        fp: False positives
-        tn: True negatives
-        fn: False negatives
-        beta: The parameter `beta` (which determines the weight of recall in the combined score)
-        ignore_index: Integer specifying a target class to ignore. If given, this class index does not contribute
-            to the returned score, regardless of reduction method
-        average: Defines the reduction that is applied
-        mdmc_average: Defines how averaging is done for multi-dimensional multi-class inputs (on top of the
-            ``average`` parameter)
-    """
-
-    return fbeta_score_ipu(
-        preds,
-        target,
-        beta=beta,
-        average=average,
-        mdmc_average=mdmc_average,
-        ignore_index=ignore_index,
-        num_classes=num_classes,
-        threshold=threshold,
-        top_k=top_k,
-        multiclass=multiclass,
-    )
-
-
-def mean_squared_error_ipu(preds: Tensor, target: Tensor, squared: bool) -> Tensor:
-    """Computes mean squared error.
-
-    Handles NaNs without reshaping tensors in order to work on IPU.
-
-    Args:
-        preds: estimated labels
-        target: ground truth labels
-        squared: returns RMSE value if set to False
-
-    Return:
-        Tensor with MSE
-    """
-    target = target.clone()
-    preds = preds.clone()
-
-    # Replace the nan-targets in the preds/target tensors by 0
-    nan_targets = target.isnan()
-    preds[nan_targets] = 0.0
-    target[nan_targets] = 0.0
-
-    # Compute the loss, and rescale by the number of nan elements
-    loss = mean_squared_error(preds, target, squared)
-
-    if squared:
-        factor = nan_targets.numel() / ((~nan_targets).sum())
-    else:
-        factor = (nan_targets.numel() / ((~nan_targets).sum())).sqrt()
-
-    loss = loss * factor
-
-    return loss
-
-
-def mean_absolute_error_ipu(preds: Tensor, target: Tensor) -> Tensor:
-    """Computes mean absolute error.
-
-    Handles NaNs without reshaping tensors in order to work on IPU.
-
-    Args:
-        preds: estimated labels
-        target: ground truth labels
-
-    Return:
-        Tensor with MAE
-    """
-    target = target.clone()
-    preds = preds.clone()
-
-    # Replace the nan-targets in the preds/target tensors by 0
-    nan_targets = target.isnan()
-    preds[nan_targets] = 0.0
-    target[nan_targets] = 0.0
-
-    # Compute the loss, and rescale by the number of nan elements
-    loss = mean_absolute_error(preds, target)
-    loss = loss * nan_targets.numel() / ((~nan_targets).sum())
-
-    return loss
diff --git a/graphium/ipu/ipu_wrapper.py b/graphium/ipu/ipu_wrapper.py
index 0ac04b883..305326015 100644
--- a/graphium/ipu/ipu_wrapper.py
+++ b/graphium/ipu/ipu_wrapper.py
@@ -122,7 +122,7 @@ def training_step(self, batch, batch_idx) -> Dict[str, Any]:
         features, labels = batch["features"], batch["labels"]
         features, labels = self.squeeze_input_dims(features, labels)
         dict_input = {"features": features, "labels": labels}
-        step_dict = super().training_step(dict_input, to_cpu=False)
+        step_dict = super().training_step(dict_input)
 
         loss = step_dict.pop("loss")
         step_dict["loss"] = self.poptorch.identity_loss(loss, reduction="mean")
@@ -132,7 +132,7 @@ def validation_step(self, batch, batch_idx) -> Dict[str, Any]:
         features, labels = batch["features"], batch["labels"]
         features, labels = self.squeeze_input_dims(features, labels)
         dict_input = {"features": features, "labels": labels}
-        step_dict = super().validation_step(dict_input, to_cpu=False)
+        step_dict = super().validation_step(dict_input)
 
         return step_dict
 
@@ -141,14 +141,14 @@ def test_step(self, batch, batch_idx) -> Dict[str, Any]:
         features, labels = batch["features"], batch["labels"]
         features, labels = self.squeeze_input_dims(features, labels)
         dict_input = {"features": features, "labels": labels}
-        step_dict = super().test_step(dict_input, to_cpu=False)
+        step_dict = super().test_step(dict_input)
 
         return step_dict
 
     def predict_step(self, **inputs) -> Dict[str, Any]:
         # Build a dictionary from the tuples
         dict_input = inputs
-        step_dict = super().predict_step(dict_input, to_cpu=False)
+        step_dict = super().predict_step(dict_input)
 
         return step_dict
 
diff --git a/graphium/trainer/__init__.py b/graphium/trainer/__init__.py
index ed2cbf2a4..1e1682e2f 100644
--- a/graphium/trainer/__init__.py
+++ b/graphium/trainer/__init__.py
@@ -2,3 +2,6 @@
 from . import metrics
 
 from .predictor import PredictorModule
+from .predictor_summaries import SingleTaskSummary
+from .predictor_summaries import MultiTaskSummary
+from .progress_bar import ProgressBarMetrics
diff --git a/graphium/trainer/metrics.py b/graphium/trainer/metrics.py
index 22361faa6..7a20e0049 100644
--- a/graphium/trainer/metrics.py
+++ b/graphium/trainer/metrics.py
@@ -12,16 +12,22 @@
 """
 
 
-from typing import Union, Callable, Optional, Dict, Any
+from typing import Union, Callable, Optional, Dict, Any, Literal, List, Tuple
 
 import sys
 
 import torch
 from torch import Tensor
+import torch.distributed as dist
 import operator as op
+from copy import deepcopy
+from loguru import logger
 
+from torch.nn.modules.loss import _Loss
 from torchmetrics.utilities.distributed import reduce
 import torchmetrics.functional.regression.mae
+from torchmetrics.utilities.data import dim_zero_cat
+from torchmetrics import Metric
 
 from graphium.utils.tensor import nan_mean
 
@@ -38,7 +44,7 @@ class Thresholder:
     def __init__(
         self,
         threshold: float,
-        operator: Union[str, Callable] = "greater",
+        operator: Union[Literal["greater", "gt", ">", "lower", "lt", "<"], Callable] = "greater",
         th_on_preds: bool = True,
         th_on_target: bool = False,
     ):
@@ -74,10 +80,10 @@ def _get_operator(operator):
         """Operator can either be a string, or a callable"""
         if isinstance(operator, str):
             op_name = operator.lower()
-            if op_name in ["greater", "gt"]:
+            if op_name in ["greater", "gt", ">"]:
                 op_str = ">"
                 operator = op.gt
-            elif op_name in ["lower", "lt"]:
+            elif op_name in ["lower", "lt", "<"]:
                 op_str = "<"
                 operator = op.lt
             else:
@@ -129,6 +135,28 @@ def __eq__(self, obj) -> bool:
         return all(is_eq)
 
 
+def _filter_nans(preds: Tensor, target: Tensor, target_nan_mask: Union[Literal[None, "none", "ignore"], int]) -> Tuple[Tensor, Tensor]:
+    """Handle the NaNs according to the chosen options"""
+
+    if target_nan_mask is None: # No NaN handling
+        return preds, target
+    
+    if target.dtype in [torch.int, torch.int16, torch.int32, torch.int64, torch.int8]:
+        target_nans = (torch.iinfo(target.dtype).min == target) | (torch.iinfo(target.dtype).max == target)
+    else:
+        target_nans = torch.isnan(target)
+    if ~target_nans.any(): # No NaNs
+        return preds, target
+    elif isinstance(target_nan_mask, (int, float)): # Replace NaNs
+        target = target.clone()
+        target[target_nans] = target_nan_mask
+    elif target_nan_mask == "ignore": # Remove NaNs
+        target = target[~target_nans]
+        preds = preds[~target_nans]
+    else:
+        raise ValueError(f"Invalid option `{target_nan_mask}`")
+    return preds, target
+
 class MetricWrapper:
     r"""
     Allows to initialize a metric from a name or Callable, and initialize the
@@ -137,10 +165,10 @@ class MetricWrapper:
 
     def __init__(
         self,
-        metric: Union[str, Callable],
+        metric: Union[str, torchmetrics.Metric, torch.nn.modules.loss._Loss],
         threshold_kwargs: Optional[Dict[str, Any]] = None,
-        target_nan_mask: Optional[Union[str, int]] = None,
-        multitask_handling: Optional[str] = None,
+        target_nan_mask: Union[Literal[None, "none", "ignore"], int] = None,
+        multitask_handling: Literal[None, "none", "flatten", "mean-per-label"] = None,
         squeeze_targets: bool = False,
         target_to_int: bool = False,
         **kwargs,
@@ -187,7 +215,7 @@ def __init__(
                 Other arguments to call with the metric
         """
 
-        self.metric, self.metric_name = self._get_metric(metric)
+        metric_class, self.metric_name = self._get_metric_class(metric)
         self.thresholder = None
         if threshold_kwargs is not None:
             self.thresholder = Thresholder(**threshold_kwargs)
@@ -198,6 +226,34 @@ def __init__(
         self.target_to_int = target_to_int
         self.kwargs = kwargs
 
+        self.metric, self.kwargs = self._initialize_metric(metric_class, self.target_nan_mask, self.multitask_handling, **self.kwargs)
+
+    @staticmethod
+    def _initialize_metric(metric, target_nan_mask, multitask_handling, **kwargs):
+        r"""
+        Initialize the metric with the provided kwargs
+        """
+    
+        if not isinstance(metric, type):
+            if callable(metric):
+                metric = MetricToConcatenatedTorchMetrics(
+                    metric_fn=metric,
+                    target_nan_mask=target_nan_mask, 
+                    multitask_handling=multitask_handling, 
+                    **kwargs)
+                return metric, kwargs
+            elif all(hasattr(metric, method) for method in ["update", "compute", "reset", "to"]):
+                return metric, kwargs
+            else:
+                raise ValueError(f"metric must be a callable, or a class with 'update', 'compute', 'reset', 'to', provided: `{type(metric)}`")
+        
+        metric = metric(**kwargs)
+        if not all(hasattr(metric, method) for method in ["update", "compute", "reset", "to"]):
+            raise ValueError(f"metric must be a callable, or a class with 'update', 'compute', 'reset', 'to', provided: `{type(metric)}`")
+
+        return metric, kwargs
+
+
     @staticmethod
     def _parse_target_nan_mask(target_nan_mask):
         """
@@ -254,20 +310,35 @@ def _parse_multitask_handling(multitask_handling, target_nan_mask):
         return multitask_handling
 
     @staticmethod
-    def _get_metric(metric):
+    def _get_metric_class(metric):
         from graphium.utils.spaces import METRICS_DICT
 
         if isinstance(metric, str):
-            metric_name = metric
-            metric = METRICS_DICT[metric]
+            metric_name = MetricWrapper._ipu_metrics_name_conversion(metric)
+            metric = METRICS_DICT[metric_name]
         else:
             metric_name = None
             metric = metric
         return metric, metric_name
-
-    def compute(self, preds: Tensor, target: Tensor) -> Tensor:
+    
+    @staticmethod
+    def _ipu_metrics_name_conversion(metric, warning=True):
         r"""
-        Compute the metric, apply the thresholder if provided, and manage the NaNs
+        Convert the metric name from the removed ipu metrics to the regular torchmetrics metrics
+        """
+        metric_name = metric
+        if metric_name.endswith("_ipu"): # For backward compatibility when loading models with metrics for ipu
+            metric_name = metric_name[:-4]
+            if metric_name == "average_precision": # A previous typo in the `spaces.py`
+                metric_name = "averageprecision"
+            if warning:
+                logger.warning(f"Using the metric `{metric_name}` instead of `{metric}`")
+        return metric_name
+
+    def update(self, preds: Tensor, target: Tensor) -> Tensor:
+        r"""
+        Update the parameters of the metric, apply the thresholder if provided, and manage the NaNs.
+        See `torchmetrics.Metric.update` for more details.
         """
         if preds.ndim == 1:
             preds = preds.unsqueeze(-1)
@@ -279,8 +350,6 @@ def compute(self, preds: Tensor, target: Tensor) -> Tensor:
         if self.thresholder is not None:
             preds, target = self.thresholder(preds, target)
 
-        target_nans = torch.isnan(target)
-
         # for the classifigression task, cast predictions from
         # (batch_size, n_targets * n_brackets) to (batch_size, n_targets, n_brackets)
         # TODO: make this more flexible to the target shape in the future
@@ -290,7 +359,7 @@ def compute(self, preds: Tensor, target: Tensor) -> Tensor:
         else:
             classifigression = False
 
-        if self.multitask_handling is None:
+        if (self.multitask_handling is None):
             # In case of no multi-task handling, apply the nan filtering, then compute the metrics
             assert (
                 self.target_nan_mask != "ignore"
@@ -300,7 +369,9 @@ def compute(self, preds: Tensor, target: Tensor) -> Tensor:
                 target = target.squeeze()
             if self.target_to_int:
                 target = target.to(int)
-            metric_val = self.metric(preds, target, **self.kwargs)
+            self.metric.update(preds, target)
+
+
         elif self.multitask_handling == "flatten":
             # Flatten the tensors, apply the nan filtering, then compute the metrics
             if classifigression:
@@ -313,16 +384,26 @@ def compute(self, preds: Tensor, target: Tensor) -> Tensor:
                 target = target.squeeze()
             if self.target_to_int:
                 target = target.to(int)
-            metric_val = self.metric(preds, target, **self.kwargs)
+            self.metric.update(preds, target)
+
+        elif isinstance(self.metric, MetricToConcatenatedTorchMetrics):
+            # NaN's and multitask handling are handled by the MetricToConcatenatedTorchMetrics
+            if self.squeeze_targets:
+                target = target.squeeze()
+            if self.target_to_int:
+                target = target.to(int)
+            self.metric.update(preds, target)
+        
         elif self.multitask_handling == "mean-per-label":
             # Loop the columns (last dim) of the tensors, apply the nan filtering, compute the metrics per column, then average the metrics
-            target_list = [target[..., ii][~target_nans[..., ii]] for ii in range(target.shape[-1])]
-            # TODO: make this more flexible to the target shape in the future
+            target_list = [target[..., ii] for ii in range(target.shape[-1])]
             if classifigression:
-                preds_list = [preds[..., i, :][~target_nans[..., i]] for i in range(preds.shape[1])]
+                preds_list = [preds[..., ii, :] for ii in range(preds.shape[1])]
             else:
-                preds_list = [preds[..., ii][~target_nans[..., ii]] for ii in range(preds.shape[-1])]
-            metric_val = []
+                preds_list = [preds[..., ii] for ii in range(preds.shape[-1])]
+
+            if not isinstance(self.metric, list):
+                self.metric = [deepcopy(self.metric) for _ in range(len(target_list))]
             for ii in range(len(target_list)):
                 try:
                     this_preds, this_target = self._filter_nans(preds_list[ii], target_list[ii])
@@ -330,44 +411,78 @@ def compute(self, preds: Tensor, target: Tensor) -> Tensor:
                         this_target = this_target.squeeze()
                     if self.target_to_int:
                         this_target = this_target.to(int)
-                    metric_val.append(self.metric(this_preds, this_target, **self.kwargs))
+                    self.metric[ii].update(this_preds, this_target)
                 except:
                     pass
-            # Average the metric
-            metric_val = nan_mean(torch.stack(metric_val))
         else:
             # Wrong option
             raise ValueError(f"Invalid option `self.multitask_handling={self.multitask_handling}`")
 
-        return metric_val
+    def compute(self) -> Tensor:
+        r"""
+        Compute the metric with the method `self.compute`
+        """
+        if isinstance(self.metric, list):
+            metrics = [metric.compute() for metric in self.metric]
+            return nan_mean(torch.stack(metrics))
+
+        return self.metric.compute()
 
-    def _filter_nans(self, preds: Tensor, target: Tensor):
-        """Handle the NaNs according to the chosen options"""
-        target_nans = torch.isnan(target)
+    def update_compute(self, preds: Tensor, target: Tensor) -> Tensor:
+        r"""
+        Update the parameters of the metric, apply the thresholder if provided, and manage the NaNs.
+        Then compute the metric with the method `self.compute`
+        """
 
-        if self.target_nan_mask is None:
-            pass
-        elif isinstance(self.target_nan_mask, (int, float)):
-            target = target.clone()
-            target[torch.isnan(target)] = self.target_nan_mask
-        elif self.target_nan_mask == "ignore":
-            target = target[~target_nans]
-            preds = preds[~target_nans]
+        self.update(preds, target)
+        return self.compute()
+
+    def reset(self) -> None:
+        r"""
+        Reset the metric with the method `self.metric.reset`
+        """
+        if isinstance(self.metric, list):
+            for metric in self.metric:
+                metric.reset()
         else:
-            raise ValueError(f"Invalid option `{self.target_nan_mask}`")
-        return preds, target
+            self.metric.reset()
+
+    def to(self, device: Union[str, torch.device]) -> None:
+        r"""
+        Move the metric to the device with the method `self.metric.to`
+        """
+        if isinstance(self.metric, list):
+            for metric in self.metric:
+                metric.to(device)
+        else:
+            self.metric.to(device)
+
+    @property
+    def device(self) -> torch.device:
+        r"""
+        Return the device of the metric with the method `self.metric.device` or `self.metric[0].device`
+        """
+        if isinstance(self.metric, list):
+            return self.metric[0].device
+        return self.metric.device
+
+
+    def _filter_nans(self, preds: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]:
+        """Handle the NaNs according to the chosen options"""
+
+        return _filter_nans(preds, target, self.target_nan_mask)
 
     def __call__(self, preds: Tensor, target: Tensor) -> Tensor:
         r"""
         Compute the metric with the method `self.compute`
         """
-        return self.compute(preds, target)
+        return self.update_compute(preds, target)
 
     def __repr__(self):
         r"""
         Control how the class is printed
         """
-        full_str = f"{self.metric.__name__}"
+        full_str = f"{self.metric.__repr__()}"
         if self.thresholder is not None:
             full_str += f"({self.thresholder})"
 
@@ -405,10 +520,198 @@ def __getstate__(self):
 
     def __setstate__(self, state: dict):
         """Reload the class from pickling."""
-        state["metric"], state["metric_name"] = self._get_metric(state["metric"])
+        state["metric"], state["metric_name"] = self._get_metric_class(state["metric"])
         thresholder = state.pop("threshold_kwargs", None)
         if thresholder is not None:
             thresholder = Thresholder(**thresholder)
         state["thresholder"] = thresholder
+        state["metric"], state["at_compute_kwargs"] = self._initialize_metric(state["metric"], state["target_nan_mask"], state["multitask_handling"], **state["kwargs"])
 
         self.__dict__.update(state)
+
+class LossWrapper():
+    r"""
+    A simple wrapper to convert any metric or loss to an equivalent of `torchmetrics.Metric`
+    by adding the `update`, `compute`, and `reset` methods to make it compatible with `MetricWrapper`.
+    However, it is simply limited to computing the average of the metric over all the updates.
+    """
+
+    def __init__(self, loss):
+        self.loss = loss
+        self.scores: List[Tensor] = []
+
+    def update(self, preds: Tensor, target: Tensor):
+        self.scores.append(self.loss(preds, target))
+
+    def compute(self):
+        if len(self.scores) == 0:
+            raise ValueError("No scores to compute")
+        elif len(self.scores) == 1:
+            return self.scores[0]
+        return nan_mean(torch.stack(self.scores))
+    
+    def to(self, device: Union[str, torch.device]):
+        for ii in range(len(self.scores)):
+            self.scores[ii] = self.scores[ii].to(device)
+
+    @property
+    def device(self) -> torch.device:
+        self.loss.device
+
+    def reset(self):
+        self.scores = []
+
+
+class MetricToMeanTorchMetrics(Metric):
+    r"""
+    A simple wrapper to convert any metric or loss to an equivalent of `torchmetrics.Metric`
+    by adding the `update`, `compute`, and `reset` methods to make it compatible with `MetricWrapper`.
+
+    However, it is limited in functionality. At each `.update()`, it computes the metric and stores in a list.
+    Then at `.compute()` it returns the average of the computed metric, while ignoring NaNs.
+    """
+    scores: List[Tensor] = []
+
+    def __init__(self, metric_fn):
+        super().__init__(dist_sync_on_step=False)
+        self.metric_fn = metric_fn
+        self.add_state("scores", default=[], dist_reduce_fx="cat")
+
+    def update(self, preds: Tensor, target: Tensor):
+        self.scores.append(self.metric_fn(preds.detach(), target))
+
+    def compute(self):
+        if len(self.scores) == 0:
+            raise ValueError("No scores to compute")
+        elif len(self.scores) == 1:
+            return self.scores[0]
+        return nan_mean(torch.stack(self.scores))
+
+
+class MetricToConcatenatedTorchMetrics(Metric):
+
+    preds: List[Tensor] # Always on CPU
+    target: List[Tensor] # Always on CPU
+
+    def __init__(self, 
+                    metric_fn: Callable,
+                    target_nan_mask: Union[Literal[None, "none", "ignore"], int] = None,
+                    multitask_handling: Literal[None, "none", "flatten", "mean-per-label"] = None,
+                    **kwargs,
+                 ):
+        r"""
+            A wrapper around the `torchmetrics.Metric` to handle the saving and syncing of `preds` and `target` tensors,
+            and moving them to the CPU.
+            This is useful for certain metrics that require to save all preds and targets, such as auroc and average_precision.
+            Otherwise, if using `MetricWrapper` with the option `mean-per-label`, the `preds` and `target` would be
+            duplicated for each label, causing major memory spikes. 
+            On top of that, all preds and targets would be on the GPU, which would cause the memory to increase at every step, 
+            and potentially lead to out-of-memory before the end of the epoch.
+
+            Parameters
+            ----------
+
+            metric_fn:
+                The metric function to use. This function should take `preds` and `target` as input, and return a scalar value.
+
+            target_nan_mask:
+                - None: Do not change behaviour if there are NaNs
+
+                - int, float: Value used to replace NaNs. For example, if `target_nan_mask==0`, then
+                  all NaNs will be replaced by zeros
+
+                - 'ignore': The NaN values will be removed from the tensor before computing the metrics.
+                  Must be coupled with the `multitask_handling='flatten'` or `multitask_handling='mean-per-label'`.
+
+            multitask_handling:
+                - None: Do not process the tensor before passing it to the metric.
+                  Cannot use the option `multitask_handling=None` when `target_nan_mask=ignore`.
+                  Use either 'flatten' or 'mean-per-label'.
+
+                - 'flatten': Flatten the tensor to produce the equivalent of a single task
+
+                - 'mean-per-label': Loop all the labels columns, process them as a single task,
+                    and average the results over each task
+                  *This option might slow down the computation if there are too many labels*
+
+        """
+                 
+        super().__init__(compute_on_cpu=True, dist_sync_on_step=False, sync_on_compute=False)
+        self.metric_fn = metric_fn
+        self.target_nan_mask = target_nan_mask
+        self.multitask_handling = multitask_handling
+        self.kwargs = kwargs
+        self.add_state("preds", default=[], dist_reduce_fx="cat")
+        self.add_state("target", default=[], dist_reduce_fx="cat")
+        self._to_device_warned: bool = False
+        super().to("cpu")
+
+    def update(self, preds: Tensor, target: Tensor):
+
+        # If distributed, gather the preds and target tensors
+        if self.dist_sync_fn is not None:
+            preds_list = self.dist_sync_fn(preds, self.process_group)
+            target_list = self.dist_sync_fn(target, self.process_group)
+            preds = dim_zero_cat(preds_list)
+            target = dim_zero_cat(target_list)
+
+        # Move the tensors to the CPU after gathering them
+        self.preds.append(preds.detach().cpu())
+        self.target.append(target.cpu())
+
+    def compute(self):
+        preds = dim_zero_cat(self.preds)
+        target = dim_zero_cat(self.target)
+
+        if (self.multitask_handling is None) or (self.multitask_handling in ["none", "flatten"]):
+            preds, target = _filter_nans(preds, target, self.target_nan_mask)
+            value = self.metric_fn(preds, target,  **self.kwargs)
+
+        elif self.multitask_handling == "mean-per-label":
+            value = []
+            # Loop the columns (last dim) of the tensors, apply the nan filtering, compute the metrics per column, then average the metrics
+            target_list = [target[..., ii] for ii in range(target.shape[-1])]
+            preds_list = [preds[..., ii] for ii in range(preds.shape[-1])]
+            for ii in range(len(target_list)):
+                try:
+                    this_preds, this_target = _filter_nans(preds_list[ii], target_list[ii], self.target_nan_mask)
+                    value.append(self.metric_fn(this_preds, this_target, **self.kwargs))
+                except:
+                    pass
+            value = nan_mean(torch.stack(value))
+        else:
+            # Wrong option
+            raise ValueError(f"Invalid option `self.multitask_handling={self.multitask_handling}`")
+        return value
+    
+    def to(self, device: Union[str, torch.device]):
+        """
+        Disables the moving of the metric to another device. Stays on CPU to avoid overflow.
+        """
+        device = torch.device(device)
+        if device == torch.device("cpu"):
+            return
+        if not self._to_device_warned:
+            self._to_device_warned = True
+            logger.warning(f"{self.get_obj_name(self)}({self.get_obj_name(self.metric_fn)}) stays on `{self.device}`, won't move to `{device}`")
+        
+    @staticmethod
+    def get_obj_name(obj):
+        """
+        Returns the name of a function, class, or instance of a class.
+        
+        Parameters:
+        - obj: The object to get the name of.
+        
+        Returns:
+        - The name of the object as a string.
+        """
+        # If the object is a class or function, return its __name__
+        if hasattr(obj, '__name__'):
+            return obj.__name__
+        # If the object is an instance of a class, return its class's __name__
+        elif hasattr(obj, '__class__'):
+            return obj.__class__.__name__
+        else:
+            return str(obj)  # Fallback to converting the object to string
+
diff --git a/graphium/trainer/predictor.py b/graphium/trainer/predictor.py
index 95d2c3166..f15521268 100644
--- a/graphium/trainer/predictor.py
+++ b/graphium/trainer/predictor.py
@@ -14,7 +14,7 @@
 
 import time
 from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, Literal, Mapping
 
 import lightning
 import numpy as np
@@ -23,20 +23,20 @@
 from mup.optim import MuAdam
 from torch import Tensor, nn
 from torch_geometric.data import Batch, Data
+from torchmetrics import Metric
 
 from graphium.config.config_convert import recursive_config_reformating
 from graphium.data.datamodule import BaseDataModule
-from graphium.trainer.metrics import MetricWrapper
+from graphium.trainer.metrics import MetricWrapper, LossWrapper
 from graphium.trainer.predictor_options import (
     EvalOptions,
     FlagOptions,
     ModelOptions,
     OptimOptions,
 )
-from graphium.trainer.predictor_summaries import TaskSummaries
+from graphium.trainer.predictor_summaries import MultiTaskSummary, GradientNormMetric
 from graphium.utils import fs
 from graphium.utils.moving_average_tracker import MovingAverageTracker
-from graphium.utils.spaces import GRAPHIUM_PRETRAINED_MODELS_DICT
 from graphium.utils.tensor import dict_tensor_fp16_to_fp32
 
 
@@ -54,15 +54,12 @@ def __init__(
         scheduler_kwargs: Optional[Dict[str, Any]] = None,
         target_nan_mask: Optional[Union[str, int]] = None,
         multitask_handling: Optional[str] = None,
-        metrics: Dict[str, Callable] = None,
+        metrics: Dict[str, Dict[str, Union[Metric, "MetricWrapper"]]] = None,
         metrics_on_progress_bar: Dict[str, List[str]] = [],
         metrics_on_training_set: Optional[Dict[str, List[str]]] = None,
         flag_kwargs: Dict[str, Any] = None,
         task_norms: Optional[Dict[Callable, Any]] = None,
         metrics_every_n_train_steps: Optional[int] = None,
-        replicas: int = 1,
-        gradient_acc: int = 1,
-        global_bs: Optional[int] = 1,
     ):
         """
         The Lightning module responsible for handling the predictions, losses, metrics, optimization, etc.
@@ -139,11 +136,12 @@ def __init__(
 
         # Task-specific evalutation attributes
         self.loss_fun = {}
+        loss_names = {}
         self.metrics = {}
         self.metrics_on_progress_bar = {}
         self.metrics_on_training_set = {}
         for task in self.tasks:
-            self.loss_fun[task] = EvalOptions.parse_loss_fun(loss_fun[task])
+            loss_names[task], self.loss_fun[task] = EvalOptions.parse_loss_fun(loss_fun[task])
             self.metrics[task] = (
                 self._eval_options_dict[task].metrics
                 if self._eval_options_dict[task].metrics is not None
@@ -164,36 +162,51 @@ def __init__(
         # Set the parameters for optimizer options
         self.optim_options.set_kwargs()
 
+        # Add the loss to the metrics
+        metrics_with_loss = deepcopy(self.metrics)
+        for task in self.tasks:
+            metrics_with_loss[task][f"loss_{loss_names[task]}"] = MetricWrapper(
+                metric=LossWrapper(self.loss_fun[task]),
+                target_nan_mask=self.target_nan_mask,
+                multitask_handling=self.multitask_handling,
+            )
+        
         # Initialize the epoch summary
-        monitor = self.optim_options.scheduler_kwargs["monitor"].split("/")[0]
-        mode = self.optim_options.scheduler_kwargs["mode"]
-
-        self.task_epoch_summary = TaskSummaries(
-            task_loss_fun=self.loss_fun,
-            task_metrics=self.metrics,
-            task_metrics_on_training_set=self.metrics_on_training_set,
-            task_metrics_on_progress_bar=self.metrics_on_progress_bar,
-            monitor=monitor,
-            mode=mode,
-        )
+        self.task_epoch_summary = {
+            "train": MultiTaskSummary(
+                task_metrics=metrics_with_loss, 
+                step_name="train", 
+                task_metrics_on_progress_bar=None,
+                task_metrics_on_training_set=self.metrics_on_training_set,
+                ),
+            "val": MultiTaskSummary(
+                task_metrics=metrics_with_loss, 
+                step_name="val", 
+                task_metrics_on_progress_bar=self.metrics_on_progress_bar,
+                task_metrics_on_training_set=None,
+                ),
+            "test": MultiTaskSummary(
+                task_metrics=metrics_with_loss,
+                step_name="test",
+                task_metrics_on_progress_bar=None,
+                task_metrics_on_training_set=None,
+            ),
+        }
 
         # This helps avoid a bug when saving hparams to yaml with different dict or str formats
         self._set_hparams(recursive_config_reformating(self.hparams))
 
         # throughput estimation
-        self.mean_val_time_tracker = MovingAverageTracker()
-        self.mean_val_tput_tracker = MovingAverageTracker()
-        self.validation_step_outputs = []
-        self.test_step_outputs = []
-        self.epoch_start_time = None
+        self.mean_time_tracker = MovingAverageTracker()
+        self.mean_tput_tracker = MovingAverageTracker()
+        self.epoch_start_time = {}
 
         # Decide whether to log every step or once at the end
         # of the epoch.
         self.metrics_every_n_train_steps = metrics_every_n_train_steps
         # Wether save preds and targets for each training step.
 
-        self.samples_seen = 0
-        self.global_bs = global_bs
+        self.model_grad = GradientNormMetric()
 
     def forward(
         self, inputs: Dict
@@ -234,6 +247,22 @@ def _get_task_key(self, task_level: str, task: str):
         if not task.startswith(task_prefix):
             task = task_prefix + task
         return task
+    
+    def _get_average_loss_from_outputs(self, outputs: Dict[Literal["loss", "task_losses"], Tensor], step_name: Literal["train", "val", "test"]) -> Dict[str, Tensor]:
+        r"""
+        Averages the loss over the different tasks
+        """
+        global_loss = torch.as_tensor(outputs["loss"]).detach()
+        if global_loss.numel() > 1:
+            global_loss = global_loss[global_loss != 0].mean()
+        average_losses = {f"_global/loss/{step_name}": global_loss}
+        for task in self.tasks:
+            this_losses = torch.as_tensor(outputs["task_losses"][task]).detach()
+            if this_losses.numel() > 1:
+                this_losses = this_losses[this_losses != 0].mean()
+            average_losses[f"{task}/loss/{step_name}"] = this_losses
+        return average_losses
+
 
     def configure_optimizers(self, impl=None):
         if impl is None:
@@ -306,7 +335,7 @@ def compute_loss(
 
         wrapped_loss_fun_dict = {
             task: MetricWrapper(
-                metric=loss,
+                metric=LossWrapper(loss),
                 threshold_kwargs=None,
                 target_nan_mask=target_nan_mask,
                 multitask_handling=multitask_handling,
@@ -316,16 +345,18 @@ def compute_loss(
 
         if weights is not None:
             raise NotImplementedError("Weights are no longer supported in the loss")
+
         all_task_losses = {
-            task: wrapped(preds=preds[task], target=targets[task])
+            task: wrapped.update_compute(preds=preds[task], target=targets[task])
             for task, wrapped in wrapped_loss_fun_dict.items()
         }
+
         total_loss = torch.sum(torch.stack(list(all_task_losses.values())), dim=0)
         num_tasks = len(all_task_losses.keys())
         weighted_loss = total_loss / num_tasks
         return weighted_loss, all_task_losses
 
-    def _general_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool) -> Dict[str, Any]:
+    def _general_step(self, batch: Dict[str, Tensor], step_name: Literal["train", "val", "test"]) -> Dict[str, Any]:
         r"""Common code for training_step, validation_step and testing_step"""
         preds = self.forward(batch)  # The dictionary of predictions
 
@@ -366,7 +397,6 @@ def _general_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool)
             multitask_handling=self.multitask_handling,
         )
 
-        device = "cpu" if to_cpu else None
         for task in preds:
             task_specific_norm = self.task_norms[task] if self.task_norms is not None else None
             if hasattr(task_specific_norm, "normalize_val_test"):
@@ -379,28 +409,18 @@ def _general_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool)
                 # if normalize_val_test is true, no denormalization is applied, all losses and metrics are normalized version
                 preds[task] = task_specific_norm.denormalize(preds[task])
                 targets_dict[task] = task_specific_norm.denormalize(targets_dict[task])
-            preds[task] = preds[task].detach().to(device=device)
-            targets_dict[task] = targets_dict[task].detach().to(device=device)
-        if weights is not None:
-            weights = weights.detach().to(device=device)
+            preds[task] = preds[task].detach()
+            targets_dict[task] = targets_dict[task].detach()
 
-        step_dict = {"preds": preds, "targets": targets_dict, "weights": weights}
-        # step_dict[f"{self.loss_fun._get_name()}/{step_name}"] = loss.detach().cpu()            original
-
-        # step_dict[f"weighted_loss/{step_name}"] = loss.detach().cpu()
-        # step_dict[f"loss/{step_name}"] = loss.detach().cpu()
-        for task in self.tasks:
-            step_dict[
-                self.task_epoch_summary.metric_log_name(task, self.loss_fun[task]._get_name(), step_name)
-            ] = loss.detach()
+        self.task_epoch_summary[step_name].update(preds, targets_dict)
 
+        step_dict = {}
         step_dict["loss"] = loss
-        # print("loss ", self.global_step, self.current_epoch, loss)
         step_dict["task_losses"] = task_losses
-        step_dict["gradient_norm"] = self.get_gradient_norm()
         return step_dict
 
-    def flag_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool) -> Dict[str, Any]:
+
+    def flag_step(self, batch: Dict[str, Tensor], step_name: Literal["train", "val", "test"]) -> Dict[str, Any]:
         r"""
         Perform adversarial data agumentation during one training step using FLAG.
         Paper: https://arxiv.org/abs/2010.09891
@@ -456,56 +476,51 @@ def flag_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool) -> D
             )
             loss = loss / n_steps
 
-        device = "cpu" if to_cpu else None
         for key in preds.keys():
-            preds[key] = preds[key].detach().to(device=device)
-            targets[key] = targets[key].detach().to(device=device)
+            preds[key] = preds[key].detach()
+            targets[key] = targets[key].detach()
         if weights is not None:
-            weights = weights.detach().to(device=device)
+            weights = weights.detach()
 
-        step_dict = {"preds": preds, "targets": targets, "weights": weights}
+        step_dict = {}
         step_dict[f"loss/{step_name}"] = loss.detach().cpu()
         step_dict["loss"] = loss
         step_dict["task_losses"] = task_losses
+        self.task_epoch_summary[step_name].update(preds, targets)
         return step_dict
 
     def on_train_batch_start(self, batch: Any, batch_idx: int) -> Optional[int]:
-        self.train_batch_start_time = time.time()
+
+        self.model_grad.reset()
+        self.task_epoch_summary["train"].reset()
+        self.batch_start_time = time.time()
         self.skip_log_train_metrics = (self.metrics_every_n_train_steps is None) or (
             (batch_idx % self.metrics_every_n_train_steps) != 0
         )
         return super().on_train_batch_start(batch, batch_idx)
 
     def on_train_batch_end(self, outputs, batch: Any, batch_idx: int) -> None:
-        train_batch_time = time.time() - self.train_batch_start_time  # To be used for throughput calculation
+        train_batch_time = time.time() - self.batch_start_time  # To be used for throughput calculation
 
         # Get the metrics that are logged at every step (loss, grad_norm, batch_time, batch_tput)
-        concatenated_metrics_logs = {}
-        concatenated_metrics_logs["train/loss"] = outputs["loss"]
-        concatenated_metrics_logs["epoch_count"] = self.current_epoch
-        # Incriment by the batch size
-        self.samples_seen += self.global_bs
-        concatenated_metrics_logs["samples_seen"] = self.samples_seen
+        metrics_logs = {}
 
         # report the training loss for each individual tasks
-        for task in self.tasks:
-            concatenated_metrics_logs[f"train/loss/{task}"] = outputs["task_losses"][task]
-
         # get the mean loss value for individual tasks as they are a tensor of size --> gradient accumulation * replication * device_iter
         # filter zeros out for the individual losses
-        for key in concatenated_metrics_logs:
-            if isinstance(concatenated_metrics_logs[key], torch.Tensor):
-                if concatenated_metrics_logs[key].numel() > 1:
-                    concatenated_metrics_logs[key] = concatenated_metrics_logs[key][
-                        concatenated_metrics_logs[key] != 0
-                    ].mean()
+        losses = self._get_average_loss_from_outputs(outputs, step_name="train")
+
+        metrics_logs.update(losses)
+        metrics_logs.update(self.task_epoch_summary["train"].compute())
 
         # If logging is skipped for this step, then log the important metrics anyway and return
         if self.skip_log_train_metrics:
-            if self.logger is not None:
-                self.logger.log_metrics(
-                    concatenated_metrics_logs, step=self.global_step
-                )  # This is a pytorch lightning function call
+            self.log_dict(
+                dictionary=metrics_logs,
+                logger=True,
+                on_step=True,
+                prog_bar=True,
+            )
             return
 
         ### The code below is not executed if the logging is skipped for this step ###
@@ -513,165 +528,130 @@ def on_train_batch_end(self, outputs, batch: Any, batch_idx: int) -> None:
         # Get the throughput of the batch
         num_graphs = self.get_num_graphs(batch["features"])
         tput = num_graphs / train_batch_time
-        concatenated_metrics_logs["train/batch_time"] = train_batch_time
-        concatenated_metrics_logs["train/batch_tput"] = tput
-
-        # Compute all the metrics for the training set
-        self.task_epoch_summary.update_predictor_state(
-            step_name="train",
-            targets=outputs["targets"],
-            preds=outputs["preds"],
-            loss=outputs["loss"],  # This is the weighted loss for now, but change to task-specific loss
-            task_losses=outputs["task_losses"],
-            n_epochs=self.current_epoch,
-        )
-        metrics_logs = self.task_epoch_summary.get_metrics_logs()  # Dict[task, metric_logs]
-        metrics_logs["_global"]["grad_norm"] = self.get_gradient_norm()
-        concatenated_metrics_logs.update(metrics_logs)
+        metrics_logs["_global/batch_time/train"] = train_batch_time
+        metrics_logs["_global/batch_tput/train"] = tput
+        self.mean_time_tracker.update(train_batch_time)
+        self.mean_tput_tracker.update(tput)
+
+        metrics_computed = self.task_epoch_summary["train"].compute()
+        self.task_epoch_summary["train"].reset()
+        metrics_logs.update(metrics_computed)
+        metrics_logs["_global/grad_norm/train"] = self.model_grad.compute()
+        self.model_grad.reset()
 
         # Log the metrics
-        if self.logger is not None:
-            self.logger.log_metrics(
-                concatenated_metrics_logs, step=self.global_step
-            )  # This is a pytorch lightning function call
+        self.log_dict(
+            dictionary=metrics_logs,
+            logger=True,
+            on_step=True,
+            prog_bar=True,
+        )
 
-    def training_step(self, batch: Dict[str, Tensor], to_cpu: bool = True) -> Dict[str, Any]:
+    def training_step(self, batch: Dict[str, Tensor]) -> Dict[str, Any]:
         step_dict = None
 
         # Train using FLAG
         if self.flag_kwargs["n_steps"] > 0:
-            step_dict = self.flag_step(batch=batch, step_name="train", to_cpu=to_cpu)
+            step_dict = self.flag_step(batch=batch, step_name="train")
         # Train normally, without using FLAG
         elif self.flag_kwargs["n_steps"] == 0:
-            # step_dict = self._general_step(batch=batch, step_name="train", to_cpu=True)
-            step_dict = self._general_step(batch=batch, step_name="train", to_cpu=to_cpu)
+            # step_dict = self._general_step(batch=batch, step_name="train")
+            step_dict = self._general_step(batch=batch, step_name="train")
+
+        # Update the gradients
+        self.model_grad.update(self.model)
 
-        # Remove the preds and targets if no logging is required
-        if self.skip_log_train_metrics:
-            step_dict.pop("preds")
-            step_dict.pop("targets")
         return step_dict  # Returning the metrics_logs with the loss
 
-    def get_gradient_norm(self):
-        # compute the norm
-        total_norm = torch.tensor(0.0)
-        for p in self.parameters():
-            if p.grad is not None:
-                param_norm = p.grad.detach().data.norm(2)
-                total_norm += param_norm.detach().cpu() ** 2
-        total_norm = total_norm**0.5
-        return total_norm
+    def validation_step(self, batch: Dict[str, Tensor]) -> Dict[str, Any]:
+        return self._general_step(batch=batch, step_name="val")
 
+    def test_step(self, batch: Dict[str, Tensor]) -> Dict[str, Any]:
+        return self._general_step(batch=batch, step_name="test")
+    
+    def _general_epoch_start(self, step_name: Literal["train", "val", "test"]) -> None:
+        self.task_epoch_summary[step_name].reset()
+        self.epoch_start_time[step_name] = time.time()
+        self.mean_time_tracker.reset()
+        self.mean_tput_tracker.reset()
     
     def predict_step(self, batch: Dict[str, Tensor]) -> Dict[str, Any]:
         preds = self.forward(batch)  # The dictionary of predictions
         targets_dict = batch.get("labels")
 
         return preds, targets_dict
-    
-    def validation_step(self, batch: Dict[str, Tensor], to_cpu: bool = True) -> Dict[str, Any]:
-        return self._general_step(batch=batch, step_name="val", to_cpu=to_cpu)
 
-    def test_step(self, batch: Dict[str, Tensor], to_cpu: bool = True) -> Dict[str, Any]:
-        return self._general_step(batch=batch, step_name="test", to_cpu=to_cpu)
 
-    def _general_epoch_end(self, outputs: Dict[str, Any], step_name: str, device: str) -> None:
+    def _general_epoch_end(self, step_name: Literal["train", "val", "test"]) -> Dict[str, Tensor]:
         r"""Common code for training_epoch_end, validation_epoch_end and testing_epoch_end"""
         # Transform the list of dict of dict, into a dict of list of dict
-        preds = {}
-        targets = {}
-        for task in self.tasks:
-            preds[task] = torch.cat([out["preds"][task].to(device) for out in outputs], dim=0)
-            targets[task] = torch.cat([out["targets"][task].to(device) for out in outputs], dim=0)
-        if ("weights" in outputs[0].keys()) and (outputs[0]["weights"] is not None):
-            weights = torch.cat([out["weights"].to(device) for out in outputs], dim=0)
-        else:
-            weights = None
+        
+        metric_logs = self.task_epoch_summary[step_name].compute()
+        self.task_epoch_summary[step_name].reset()
+        metric_logs_cpu = {k: v for k, v in metric_logs.items() if v.device == torch.device("cpu")}
+        if len(metric_logs_cpu) > 0:
+            self.log_dict(metric_logs_cpu, logger=True, prog_bar=True, sync_dist=False, on_epoch=True)
+        
+        metric_logs_accelerator = {k: v for k, v in metric_logs.items() if v.device != torch.device("cpu")}
+        if len(metric_logs_accelerator) > 0:
+            self.log_dict(metric_logs_accelerator, logger=True, prog_bar=True, sync_dist=True, on_epoch=True)
 
-        # NOTE: Computing the loss over the entire split may cause
-        # overflow issues when using fp16
-        loss, task_losses = self.compute_loss(
-            preds=dict_tensor_fp16_to_fp32(preds),
-            targets=dict_tensor_fp16_to_fp32(targets),
-            weights=weights,
-            target_nan_mask=self.target_nan_mask,
-            multitask_handling=self.multitask_handling,
-            loss_fun=self.loss_fun,
-        )
+        # Time metrics are tracked always on CPU, without progress bar, so we log them separatly
+        time_metrics = {}
+        time_metrics[f"_global/mean_batch_time/{step_name}"] = torch.tensor(self.mean_time_tracker.mean_value)
+        time_metrics[f"_global/mean_tput/{step_name}"] = self.mean_tput_tracker.mean_value
+        time_metrics[f"_global/epoch_time/{step_name}"] = torch.tensor(time.time() - self.epoch_start_time[step_name])
 
-        self.task_epoch_summary.update_predictor_state(
-            step_name=step_name,
-            preds=preds,
-            targets=targets,
-            loss=loss,
-            task_losses=task_losses,
-            n_epochs=self.current_epoch,
-        )
-        metrics_logs = self.task_epoch_summary.get_metrics_logs()
-        
-        for task in metrics_logs.keys():
-            for key, val in metrics_logs[task].items():
-                metrics_logs[task][key] = val.to(self.device)
+        self.log_dict(time_metrics, logger=True, prog_bar=False, sync_dist=False, on_epoch=True)
 
-        return metrics_logs  # Consider returning concatenated dict for logging
+        return metric_logs
 
     def on_train_epoch_start(self) -> None:
-        self.epoch_start_time = time.time()
+        self._general_epoch_start(step_name="train")
 
     def on_train_epoch_end(self) -> None:
-        if self.epoch_start_time is None:
-            logger.warning("epoch timer not initialized")
-        else:
-            epoch_time = time.time() - self.epoch_start_time
-            self.epoch_start_time = None
-            self.log("epoch_time", torch.tensor(epoch_time).to(self.device), sync_dist=True)
+        self._general_epoch_end(step_name="train")
 
     def on_validation_epoch_start(self) -> None:
-        self.mean_val_time_tracker.reset()
-        self.mean_val_tput_tracker.reset()
+        self._general_epoch_start(step_name="val")
         return super().on_validation_epoch_start()
 
     def on_validation_batch_start(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
-        self.validation_batch_start_time = time.time()
+        self.batch_start_time = time.time()
         return super().on_validation_batch_start(batch, batch_idx, dataloader_idx)
 
     def on_validation_batch_end(
-        self, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int = 0
+        self, outputs, batch: Any, batch_idx: int, dataloader_idx: int = 0
     ) -> None:
-        val_batch_time = time.time() - self.validation_batch_start_time
-        self.validation_step_outputs.append(outputs)
-        self.mean_val_time_tracker.update(val_batch_time)
+        val_batch_time = time.time() - self.batch_start_time
+        self.mean_time_tracker.update(val_batch_time)
         num_graphs = self.get_num_graphs(batch["features"])
-        self.mean_val_tput_tracker.update(num_graphs / val_batch_time)
+        self.mean_tput_tracker.update(num_graphs / val_batch_time)
         return super().on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx)
 
     def on_validation_epoch_end(self) -> None:
-        metrics_logs = self._general_epoch_end(
-            outputs=self.validation_step_outputs, step_name="val", device="cpu"
-        )
-        self.validation_step_outputs.clear()
-        concatenated_metrics_logs = self.task_epoch_summary.concatenate_metrics_logs(metrics_logs, device=self.device)
-        concatenated_metrics_logs["val/mean_time"] = torch.tensor(self.mean_val_time_tracker.mean_value).to(self.device)
-        concatenated_metrics_logs["val/mean_tput"] = self.mean_val_tput_tracker.mean_value.to(self.device)
-        self.log_dict(concatenated_metrics_logs, sync_dist=True)
+        self._general_epoch_end(step_name="val")
+        return super().on_validation_epoch_end()
 
-        # Save yaml file with the per-task metrics summaries
-        full_dict = {}
-        full_dict.update(self.task_epoch_summary.get_dict_summary())
-
-    def on_test_batch_end(self, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
-        self.test_step_outputs.append(outputs)
+    def on_test_epoch_start(self) -> None:
+        self._general_epoch_start(step_name="test")
+        return super().on_test_epoch_start()
 
     def on_test_epoch_end(self) -> None:
-        metrics_logs = self._general_epoch_end(outputs=self.test_step_outputs, step_name="test", device="cpu")
-        self.test_step_outputs.clear()
-        concatenated_metrics_logs = self.task_epoch_summary.concatenate_metrics_logs(metrics_logs, device=self.device)
-
-        self.log_dict(concatenated_metrics_logs, sync_dist=True)
 
-        # Save yaml file with the per-task metrics summaries
-        full_dict = {}
-        full_dict.update(self.task_epoch_summary.get_dict_summary())
+        self._general_epoch_end(step_name="test")
+        return super().on_test_epoch_end()
+    
+    def on_test_batch_start(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
+        self.batch_start_time = time.time()
+        return super().on_test_batch_start(batch, batch_idx, dataloader_idx)
+    
+    def on_test_batch_end(self, outputs: Tensor | Mapping[str, Any] | None, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
+        test_batch_time = time.time() - self.batch_start_time
+        self.mean_time_tracker.update(test_batch_time)
+        num_graphs = self.get_num_graphs(batch["features"])
+        self.mean_tput_tracker.update(num_graphs / test_batch_time)
+        return super().on_test_batch_end(outputs, batch, batch_idx, dataloader_idx)
 
     def on_train_start(self):
         hparams_log = deepcopy(self.hparams)
@@ -679,16 +659,15 @@ def on_train_start(self):
         if self.logger is not None:
             self.logger.log_hyperparams(hparams_log)
 
-    def get_progress_bar_dict(self) -> Dict[str, float]:
-        prog_dict = {}
-        prog_dict["loss"] = self.task_epoch_summary.weighted_loss.detach().cpu()
-        results_on_progress_bar = self.task_epoch_summary.get_results_on_progress_bar("val")
-        for task in self.tasks:
-            prog_dict[self.task_epoch_summary.metric_log_name(task, "loss", "val")] = (
-                self.task_epoch_summary.task_summaries[task].summaries["val"].loss
-            )
-            prog_dict.update(results_on_progress_bar)
-        return prog_dict
+    @property
+    def get_metrics_on_progress_bar(self) -> List[str]:
+        prog_list = ["_global/loss/train"]
+        for task_name in self.tasks:
+            for metric in self.metrics_on_progress_bar[task_name]:
+                this_summary = self.task_epoch_summary["val"][task_name]
+                prog_list.append(this_summary.metric_log_name(metric))
+
+        return prog_list
 
     def __repr__(self) -> str:
         r"""
@@ -702,10 +681,12 @@ def __repr__(self) -> str:
     @staticmethod
     def list_pretrained_models():
         """List available pretrained models."""
-        return GRAPHIUM_PRETRAINED_MODELS_DICT
+        from graphium.utils.spaces import GRAPHIUM_PRETRAINED_MODELS_DICT
+
+        return GRAPHIUM_PRETRAINED_MODELS_DICT # Avoiding circular imports with `space.py`
 
     @staticmethod
-    def load_pretrained_model(name_or_path: str, device: str = None):
+    def load_pretrained_model(name_or_path: str, device: str = None, strict: bool = True, **kwargs):
         """Load a pretrained model from its name.
 
         Args:
@@ -713,11 +694,13 @@ def load_pretrained_model(name_or_path: str, device: str = None):
                 from `graphium.trainer.PredictorModule.list_pretrained_models()`.
         """
 
+        from graphium.utils.spaces import GRAPHIUM_PRETRAINED_MODELS_DICT # Avoiding circular imports with `space.py`
+
         name = GRAPHIUM_PRETRAINED_MODELS_DICT.get(name_or_path)
 
         if name is not None:
             return PredictorModule.load_from_checkpoint(
-                GRAPHIUM_PRETRAINED_MODELS_DICT[name_or_path], map_location=device
+                GRAPHIUM_PRETRAINED_MODELS_DICT[name_or_path], map_location=device, strict=strict, **kwargs
             )
 
         if name is None and not (fs.exists(name_or_path) and fs.get_extension(name_or_path) == "ckpt"):
@@ -726,7 +709,7 @@ def load_pretrained_model(name_or_path: str, device: str = None):
                 "or pass a valid checkpoint (.ckpt) path."
             )
 
-        return PredictorModule.load_from_checkpoint(name_or_path, map_location=device)
+        return PredictorModule.load_from_checkpoint(name_or_path, map_location=device, strict=strict, **kwargs)
 
     def set_max_nodes_edges_per_graph(self, datamodule: BaseDataModule, stages: Optional[List[str]] = None):
         datamodule.setup()
diff --git a/graphium/trainer/predictor_options.py b/graphium/trainer/predictor_options.py
index 04a62e84b..5303976a9 100644
--- a/graphium/trainer/predictor_options.py
+++ b/graphium/trainer/predictor_options.py
@@ -30,10 +30,6 @@
 
 from torch import nn
 
-from graphium.utils.spaces import LOSS_DICT
-from graphium.utils.spaces import SCHEDULER_DICT
-
-
 @dataclass
 class ModelOptions:
     r"""
@@ -117,6 +113,7 @@ def set_kwargs(self):
         scheduler_class = torch_scheduler_kwargs.pop("module_type")
         if self.scheduler_class is None:
             if isinstance(scheduler_class, str):
+                from graphium.utils.spaces import SCHEDULER_DICT
                 self.scheduler_class = SCHEDULER_DICT[scheduler_class]
             elif isclass(scheduler_class):
                 self.scheduler_class = scheduler_class
@@ -196,12 +193,15 @@ def parse_loss_fun(loss_fun: Union[str, Dict, Callable]) -> Callable:
                 Function or callable to compute the loss, takes `preds` and `targets` as inputs.
         """
 
+        from graphium.utils.spaces import LOSS_DICT # Avoiding circular imports with `spaces.py`
+
         if isinstance(loss_fun, str):
             if loss_fun not in LOSS_DICT.keys():
                 raise ValueError(
                     f"`loss_fun` expected to be one of the strings in {LOSS_DICT.keys()}. "
                     f"Provided: {loss_fun}."
                 )
+            loss_name = loss_fun
             loss_fun = LOSS_DICT[loss_fun]()
         elif isinstance(loss_fun, dict):
             if loss_fun.get("name") is None:
@@ -214,10 +214,12 @@ def parse_loss_fun(loss_fun: Union[str, Dict, Callable]) -> Callable:
             loss_fun = deepcopy(loss_fun)
             loss_name = loss_fun.pop("name")
             loss_fun = LOSS_DICT[loss_name](**loss_fun)
-        elif not callable(loss_fun):
+        elif callable(loss_fun):
+            loss_name = str(loss_fun)
+        else:
             raise ValueError(f"`loss_fun` must be `str`, `dict` or `callable`. Provided: {type(loss_fun)}")
 
-        return loss_fun
+        return loss_name, loss_fun
 
 
 @dataclass
diff --git a/graphium/trainer/predictor_summaries.py b/graphium/trainer/predictor_summaries.py
index 7a4d9292f..06e762dde 100644
--- a/graphium/trainer/predictor_summaries.py
+++ b/graphium/trainer/predictor_summaries.py
@@ -14,14 +14,18 @@
 
 r"""Classes to store information about resulting evaluation metrics when using a Predictor Module."""
 
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union, Literal, Iterable, Set
 from loguru import logger
+from copy import deepcopy
+import inspect
 
 import numpy as np
 import torch
 from torch import Tensor
+from torchmetrics import MeanMetric, Metric
+from torchmetrics.aggregation import BaseAggregator
 
-from graphium.utils.tensor import nan_mean, nan_std, nan_median, tensor_fp16_to_fp32
+from graphium.trainer.metrics import MetricToConcatenatedTorchMetrics
 
 
 class SummaryInterface(object):
@@ -29,37 +33,30 @@ class SummaryInterface(object):
     An interface to define the functions implemented by summary classes that implement SummaryInterface.
     """
 
-    def set_results(self, **kwargs):
+    def update(self, preds: Tensor, targets: Tensor) -> None:
         raise NotImplementedError()
 
-    def get_dict_summary(self):
+    def compute(self, **kwargs) -> Tensor:
         raise NotImplementedError()
-
-    def update_predictor_state(self, **kwargs):
-        raise NotImplementedError()
-
-    def get_metrics_logs(self, **kwargs):
+    
+    def reset(self) -> None:
         raise NotImplementedError()
 
 
-class Summary(SummaryInterface):
-    # TODO (Gabriela): Default argument cannot be []
+class SingleTaskSummary(SummaryInterface):
     def __init__(
         self,
-        loss_fun: Union[str, Callable],
-        metrics: Dict[str, Callable],
-        metrics_on_training_set: List[str] = [],
-        metrics_on_progress_bar: List[str] = [],
-        monitor: str = "loss",
-        mode: str = "min",
+        metrics: Dict[str, Union[Metric, "MetricWrapper"]],
+        step_name: str,
+        metrics_on_training_set: Optional[List[str]] = None,
+        metrics_on_progress_bar: Optional[List[str]] = None,
         task_name: Optional[str] = None,
+        compute_mean: bool = True,
+        compute_std: bool = True,
     ):
         r"""
         A container to be used by the Predictor Module that stores the results for the given metrics on the predictions and targets provided.
         Parameters:
-            loss_fun:
-            Loss function used during training. Acceptable strings are 'mse', 'bce', 'mae', 'cosine'.
-            Otherwise, a callable object must be provided, with a method `loss_fun._get_name()`.
 
             metrics:
             A dictionnary of metrics to compute on the prediction, other than the loss function.
@@ -67,394 +64,293 @@ def __init__(
 
             metrics_on_training_set:
             The metrics names from `metrics` to be computed on the training set for each iteration.
-            If `None`, all the metrics are computed. Using less metrics can significantly improve
-            performance, depending on the number of readouts.
+            If `None`, no metrics are computed.
 
             metrics_on_progress_bar:
-            The metrics names from `metrics` to display also on the progress bar of the training
-
-            monitor:
-            `str` metric to track (Default=`"loss/val"`)
+            The metrics names from `metrics` to display also on the progress bar of the training.
+            If `None`, no metrics are displayed.
 
             task_name:
             name of the task (Default=`None`)
 
+            compute_mean:
+            whether to compute the mean of the predictions and targets
+
+            compute_std:
+            whether to compute the standard deviation of the predictions and targets
+
         """
-        self.loss_fun = loss_fun
-        self.metrics = metrics
-        self.metrics_on_training_set = metrics_on_training_set
-        self.metrics_on_progress_bar = metrics_on_progress_bar
-        self.monitor = monitor
-        self.mode = mode
+        self.step_name = step_name
+        self.compute_mean = compute_mean
+        self.compute_std = compute_std
 
-        self.summaries = {}
-        self.best_summaries = {}
+        if not isinstance(metrics, dict):
+            raise ValueError(f"metrics must be a dictionary. Got {type(metrics)}")
+        self.metrics = deepcopy(metrics)
 
         # Current predictor state
         # self.predictor_outputs = None
-        self.step_name: str = None
-        self.targets: Tensor = None
-        self.preds: Tensor = None
-        self.loss = None  # What type?
-        self.n_epochs: int = None
-
         self.task_name = task_name
-        self.logged_metrics_exceptions = []  # Track which metric exceptions have been logged
+        self.logged_metrics_exceptions: List[str] = []  # Track which metric exceptions have been logged
+        self.last_metrics_exceptions: List[str] = []  # Track which metric exceptions have been logged
+
+        # Add default metrics
+        if ("mean_preds" not in self.metrics) and compute_mean:
+            self.metrics["mean_preds"] = MeanMetric(nan_strategy="ignore")
+        if ("mean_target" not in self.metrics) and compute_mean:
+            self.metrics["mean_target"] = MeanMetric(nan_strategy="ignore")
+        if ("std_preds" not in self.metrics) and compute_std:
+            self.metrics["std_preds"] = STDMetric(nan_strategy="ignore")
+        if ("std_target" not in self.metrics) and compute_std:
+            self.metrics["std_target"] = STDMetric(nan_strategy="ignore")
+
+        # Parse the metrics filters
+        self.metrics_on_training_set = self._parse_metrics_filter(metrics_on_training_set)
+        self.metrics_on_progress_bar = self._parse_metrics_filter(metrics_on_progress_bar)
+
+        # Update the metrics to compute on the training set
+        if self.compute_mean:
+            self.metrics_on_training_set.update(["mean_preds", "mean_target"])
+        if self.compute_std:
+            self.metrics_on_training_set.update(["std_preds", "std_target"])
+
+        self._cached_metrics: Dict[str, Tensor] = {}
+        self._logged_warnings: Set[str] = set() # Set to track which metrics have been logged
+        self._device: torch.device = None
+
+    @property
+    def get_cached_metrics(self) -> Dict[str, Tensor]:
+        return deepcopy(self._cached_metrics)
+
+    def _parse_metrics_filter(self, filter: Optional[Union[List[str], Dict[str, Any]]]) -> List[str]:
+        if filter is None:
+            filter = []
+        elif isinstance(filter, dict):
+            filter = list(filter.keys())
+        elif isinstance(filter, (list, tuple, set)):
+            filter = list(filter)
+        elif isinstance(filter, str):
+            filter = [filter]
+        else:
+            raise ValueError(f"metrics_to_use must be a list or a dictionary. Got {type(filter)}")
 
-    def update_predictor_state(
-        self, step_name: str, targets: Tensor, preds: Tensor, loss: Tensor, n_epochs: int
-    ):
+        # Ensure that the filter is a subset of the metrics
+        all_metrics = set(self.metrics.keys())
+        filter = set(filter)
+        if not filter.issubset(all_metrics):
+            raise ValueError(f"metrics_to_use must be a subset of the metrics. Got {filter - all_metrics}, available {all_metrics}")
+
+        return filter
+
+    @property
+    def metrics_to_use(self) -> Dict[str, Callable]:
         r"""
-        update the state of the predictor
-        Parameters:
-            step_name: which stage you are in, e.g. "train"
-            targets: the targets tensor
-            predictions: the predictions tensor
-            loss: the loss tensor
-            n_epochs: the number of epochs
+        return the metrics to use by filtering the metrics dictionary if it is the training step. Otherwise, return all metrics.
         """
-        self.step_name = step_name
-        self.targets = targets
-        self.preds = preds
-        self.loss = loss
-        self.n_epochs = n_epochs
 
-    def set_results(
-        self,
-        metrics: Dict[str, Tensor],
-    ):
+        if self.step_name == "train":
+            metrics_to_use = {
+                key: metric for key, metric in self.metrics.items() if key in self.metrics_on_training_set
+            }
+
+            return metrics_to_use
+        return self.metrics
+    
+    @staticmethod
+    def _update(metric_key:str, metric_obj, preds: Tensor, targets: Tensor) -> None:
         r"""
-        set the reults from the metrics
-        [!] This function requires that self.update_predictor_state() be called before it.
+        update the state of the metrics
         Parameters:
-            metrics: a dictionary of metrics
+            targets: the targets tensor
+            predictions: the predictions tensor
         """
 
-        # Include the task_name in the loss for logging, and similarly for other metrics
-        metrics[self.metric_log_name(self.task_name, "loss", self.step_name)] = self.loss
-        self.summaries[self.step_name] = Summary.Results(
-            targets=self.targets,
-            preds=self.preds,
-            loss=self.loss,
-            metrics=metrics,  # Should include task name from get_metrics_logs()
-            monitored_metric=f"{self.monitor}/{self.step_name}",  # Include task name?
-            n_epochs=self.n_epochs,
-        )
-        if self.is_best_epoch(self.step_name, self.loss, metrics):
-            self.best_summaries[self.step_name] = self.summaries[self.step_name]
+        # Check the `metric_obj.update` signature to know if it takes `preds` and `targets` or only one of them
+        varnames = [val.name for val in inspect.signature(metric_obj.update).parameters.values()]
+        if ("preds" == varnames[0]) and ("target" == varnames[1]):
+            # The typical case of `torchmetrics`
+            metric_obj.update(preds, targets)
+        elif ("preds" == varnames[1]) and ("target" == varnames[0]):
+            # Unusual case where the order of the arguments is reversed
+            metric_obj.update(targets, preds)
+        elif ("value" == varnames[0]) and ("preds" in metric_key):
+            # The case where the metric takes only one value, and it is the prediction
+            metric_obj.update(preds)
+        elif ("value" == varnames[0]) and ("target" in metric_key):
+            # The case where the metric takes only one value, and it is the target
+            metric_obj.update(targets)
+        else:
+            raise ValueError(f"Metric {metric_key} update method signature `{varnames}` is not recognized.")
 
-    def is_best_epoch(self, step_name: str, loss: Tensor, metrics: Dict[str, Tensor]) -> bool:
+
+    def update(self, preds: Tensor, targets: Tensor) -> None:
         r"""
-        check if the current epoch is the best epoch based on self.mode criteria
+        update the state of the metrics
         Parameters:
-            step_name: which stage you are in, e.g. "train"
-            loss: the loss tensor
-            metrics: a dictionary of metrics
+            targets: the targets tensor
+            predictions: the predictions tensor
         """
 
-        # TODO (Gabriela): Check for bugs related to monitor_name
-        if not (step_name in self.best_summaries.keys()):
-            return True
-
-        # Include the task_name in the loss for logging, and similarly for other metrics
-        metrics[self.metric_log_name(self.task_name, "loss", self.step_name)] = loss
-        monitor_name = f"{self.monitor}/{step_name}"  # Include task_name?
-        if (
-            not monitor_name in self.best_summaries.keys()
-        ):  # Feels like there's a bug here. What is this trying to do???
-            return True
-
-        if self.mode == "max":
-            return metrics[monitor_name] > self.best_summaries[step_name].monitored
-        elif self.mode == "min":
-            return metrics[monitor_name] < self.best_summaries[step_name].monitored
+        self._device = preds.device
+
+        for metric_key, metric_obj in self.metrics_to_use.items():
+            metric_obj.to(self.device)
+            try:
+                self._update(metric_key, metric_obj, preds, targets)
+            except Exception as err:
+                err_msg = f"Error for metric {metric_key} on task {self.task_name} and step {self.step_name}. Exception: {err}"
+                # Check if the error is due to the device mismatch, cast to the device, and retry
+
+                if err_msg not in self._logged_warnings:
+                    logger.warning(err_msg)
+                    self._logged_warnings.add(err_msg)
+                
+
+    def _compute(self, metrics_to_use: Optional[Union[List[str], Dict[str, Any]]] = None) -> Dict[str, Tensor]:
+
+        # Parse the metrics to use
+        if metrics_to_use is None:
+            metrics_to_use = list(self.metrics.keys())
+        elif isinstance(metrics_to_use, dict):
+            metrics_to_use = list(metrics_to_use.keys())
         else:
-            ValueError(f"Mode must be 'min' or 'max', provided `{self.mode}`")
+            raise ValueError(f"metrics_to_use must be a list or a dictionary. Got {type(metrics_to_use)}")
+        
+        self.last_metrics_exceptions = []  # Reset the exceptions for this step
+
+        # Compute the metrics
+        computed_metrics = {}
+        for metric_key in metrics_to_use:
+            metric_name = self.metric_log_name(metric_key)
+            metric_obj = self.metrics[metric_key]
+            try:
+                computed_metrics[f"{metric_name}"] = metric_obj.compute()
+            except Exception as e:
+                # If the metric computation fails, return NaN and log a warning only once
+                computed_metrics[f"{metric_name}"] = torch.tensor(torch.nan, device=self.device)
+                # Warn only if it's the first warning for that metric
+                if metric_name not in self.logged_metrics_exceptions:
+                    self.logged_metrics_exceptions.append(metric_name)
+                    logger.warning(f"Error for metric {metric_name}. NaN is returned. Exception: {e}")
+                self.last_metrics_exceptions.append(metric_name)
 
-    def get_results(
-        self,
-        step_name: str,
-    ):
+        return computed_metrics
+
+    def compute(self) -> Dict[str, Tensor]:
         r"""
-        retrieve the results for a given step
-        Parameters:
-            step_name: which stage you are in, e.g. "train"
+        compute the metrics
         Returns:
-            the results for the given step
+            the computed metrics
         """
-        return self.summaries[step_name]
+        computed_metrics = self._compute(metrics_to_use=self.metrics_to_use)
+        self._cached_metrics = computed_metrics
 
-    def get_best_results(
-        self,
-        step_name: str,
-    ):
+        return computed_metrics
+    
+    def reset(self) -> None:
         r"""
-        retrieve the best results for a given step
-        Parameters:
-            step_name: which stage you are in, e.g. "train"
-        Returns:
-            the best results for the given step
+        reset the state of the metrics
         """
-        return self.best_summaries[step_name]
+        for metric_key, metric in self.metrics.items():
+            try:
+                metric.reset()
+            except AttributeError as e:
+                metric_name = self.metric_log_name(metric_key)
+                # Skip error if the message is `AttributeError: 'Tensor' object has no attribute 'clear'. Did you mean: 'char'?`
+                # This error happens when there's nothing to reset, usually because the metric failed.
+                if (metric_name not in self.last_metrics_exceptions) or ("'Tensor' object has no attribute 'clear'" not in str(e)):
+                    raise e
 
     def get_results_on_progress_bar(
         self,
-        step_name: str,
     ) -> Dict[str, Tensor]:
         r"""
         retrieve the results to be displayed on the progress bar for a given step
-        Parameters:
-            step_name: which stage you are in, e.g. "train"
-        Returns:
-            the results to be displayed on the progress bar for the given step
-        """
-        results = self.summaries[step_name]
-        results_prog = {
-            # f"{kk}/{step_name}": results.metrics[f"{kk}/{step_name}"] for kk in self.metrics_on_progress_bar
-            self.metric_log_name(self.task_name, kk, step_name): results.metrics[
-                self.metric_log_name(self.task_name, kk, step_name)
-            ]
-            for kk in self.metrics_on_progress_bar
-        }
-        return results_prog
 
-    def get_dict_summary(self) -> Dict[str, Any]:
-        r"""
-        retrieve the full summary in a dictionary
         Returns:
-            the full summary in a dictionary
-        """
-        full_dict = {}
-        # Get metric summaries
-        full_dict["metric_summaries"] = {}
-        for key, val in self.summaries.items():
-            full_dict["metric_summaries"][key] = {k: v for k, v in val.metrics.items()}
-            full_dict["metric_summaries"][key]["n_epochs"] = val.n_epochs
-
-        # Get metric summaries at best epoch
-        full_dict["best_epoch_metric_summaries"] = {}
-        for key, val in self.best_summaries.items():
-            full_dict["best_epoch_metric_summaries"][key] = val.metrics
-            full_dict["best_epoch_metric_summaries"][key]["n_epochs"] = val.n_epochs
-
-        return full_dict
-
-    def get_metrics_logs(self) -> Dict[str, Any]:
-        r"""
-        Get the data about metrics to log.
-        Note: This function requires that self.update_predictor_state() be called before it.
-        Returns:
-            A dictionary of metrics to log.
+            the results to be displayed on the progress bar for the given step
         """
+        cached_metrics = self.get_cached_metrics
+        if cached_metrics is None:
+            results_prog = self._compute(metrics_to_use=self.metrics_on_progress_bar)
+        else:
+            results_prog = {}
+            for metric_key in self.metrics_on_progress_bar:
+                metric_name = self.metric_log_name(metric_key)
+                results_prog[metric_name] = cached_metrics[metric_name]
 
-        targets = tensor_fp16_to_fp32(self.targets)
-        preds = tensor_fp16_to_fp32(self.preds)
-
-        targets = targets.to(dtype=preds.dtype, device=preds.device)
-
-        # Compute the metrics always used in regression tasks
-        metric_logs = {}
-        metric_logs[self.metric_log_name(self.task_name, "mean_pred", self.step_name)] = nan_mean(preds)
-        metric_logs[self.metric_log_name(self.task_name, "std_pred", self.step_name)] = nan_std(preds)
-        metric_logs[self.metric_log_name(self.task_name, "median_pred", self.step_name)] = nan_median(preds)
-        metric_logs[self.metric_log_name(self.task_name, "mean_target", self.step_name)] = nan_mean(targets)
-        metric_logs[self.metric_log_name(self.task_name, "std_target", self.step_name)] = nan_std(targets)
-        metric_logs[self.metric_log_name(self.task_name, "median_target", self.step_name)] = nan_median(
-            targets
-        )
-
-        # Specify which metrics to use
-        metrics_to_use = self.metrics
-        if self.step_name == "train":
-            metrics_to_use = {
-                key: metric for key, metric in metrics_to_use.items() if key in self.metrics_on_training_set
-            }
-        # Compute the additional metrics
-        for key, metric in metrics_to_use.items():
-            metric_name = self.metric_log_name(
-                self.task_name, key, self.step_name
-            )  # f"{key}/{self.step_name}"
-            try:
-                metric_logs[metric_name] = metric(preds, targets)
-            except Exception as e:
-                metric_logs[metric_name] = torch.as_tensor(float("nan"))
-                # Warn only if it's the first warning for that metric
-                if metric_name not in self.logged_metrics_exceptions:
-                    self.logged_metrics_exceptions.append(metric_name)
-                    logger.warning(f"Error for metric {metric_name}. NaN is returned. Exception: {e}")
-
-        # Convert all metrics to CPU, except for the loss
-        # metric_logs[f"{self.loss_fun._get_name()}/{self.step_name}"] = self.loss.detach().cpu()
-        metric_logs[
-            self.metric_log_name(self.task_name, self.loss_fun._get_name(), self.step_name)
-        ] = self.loss.detach().cpu()
-        # print("Metrics logs keys: ", metric_logs.keys())
-        metric_logs = {key: metric.detach().cpu() for key, metric in metric_logs.items()}
-
-        return metric_logs
+        return results_prog
 
-    def metric_log_name(self, task_name, metric_name, step_name):
-        if task_name is None:
-            return f"{metric_name}/{step_name}"
+    def metric_log_name(self, metric_name):
+        if self.task_name is None:
+            return f"{metric_name}/{self.step_name}"
         else:
-            return f"{task_name}/{metric_name}/{step_name}"
-
-    class Results:
-        def __init__(
-            self,
-            targets: Tensor = None,
-            preds: Tensor = None,
-            loss: float = None,  # Is this supposed to be a Tensor or float?
-            metrics: dict = None,
-            monitored_metric: str = None,
-            n_epochs: int = None,
-        ):
-            r"""
-            This inner class is used as a container for storing the results of the summary.
-            Parameters:
-                targets: the targets
-                preds: the prediction tensor
-                loss: the loss, float or tensor
-                metrics: the metrics
-                monitored_metric: the monitored metric
-                n_epochs: the number of epochs
-            """
-            self.targets = targets.detach().cpu()
-            self.preds = preds.detach().cpu()
-            self.loss = loss.item() if isinstance(loss, Tensor) else loss
-            self.monitored_metric = monitored_metric
-            if monitored_metric in metrics.keys():
-                self.monitored = metrics[monitored_metric].detach().cpu()
-            self.metrics = {
-                key: value.tolist() if isinstance(value, (Tensor, np.ndarray)) else value
-                for key, value in metrics.items()
-            }
-            self.n_epochs = n_epochs
+            return f"{self.task_name}/{metric_name}/{self.step_name}"
+        
+    @property
+    def device(self) -> Optional[torch.device]:
+        return self._device
 
 
-class TaskSummaries(SummaryInterface):
+class MultiTaskSummary(SummaryInterface):
     def __init__(
         self,
-        task_loss_fun: Callable,
-        task_metrics: Dict[str, Callable],
-        task_metrics_on_training_set: List[str],
-        task_metrics_on_progress_bar: List[str],
-        monitor: str = "loss",
-        mode: str = "min",
+        task_metrics: Dict[str, Dict[str, Union[Metric, "MetricWrapper"]]],
+        step_name: str,
+        task_metrics_on_training_set: Optional[Dict[str, List[str]]] = None,
+        task_metrics_on_progress_bar: Optional[Dict[str, List[str]]] = None,
+        compute_mean: bool = True,
+        compute_std: bool = True,
     ):
         r"""
         class to store the summaries of the tasks
         Parameters:
-            task_loss_fun: the loss function for each task
-            task_metrics: the metrics for each task
-            task_metrics_on_training_set: the metrics to use on the training set
-            task_metrics_on_progress_bar: the metrics to use on the progress bar
-            monitor: the metric to monitor
-            mode: the mode of the metric to monitor
+
+        
+            compute_mean:
+            whether to compute the mean of the predictions and targets
+
+            compute_std:
+            whether to compute the standard deviation of the predictions and targets
+
         """
-        self.task_loss_fun = task_loss_fun
         self.task_metrics = task_metrics
-        self.task_metrics_on_progress_bar = task_metrics_on_progress_bar
-        self.task_metrics_on_training_set = task_metrics_on_training_set
-        self.monitor = monitor
-        self.mode = mode
-
-        self.task_summaries: Dict[str, Summary] = {}
-        self.task_best_summaries: Dict[str, Summary] = {}
-        self.tasks = list(task_loss_fun.keys())
+        self.task_metrics_on_progress_bar = task_metrics_on_progress_bar if task_metrics_on_progress_bar is not None else {}
+        self.task_metrics_on_training_set = task_metrics_on_training_set if task_metrics_on_training_set is not None else {}
 
+        # Initialize all the single-task summaries
+        self.tasks = list(task_metrics.keys())
+        self.task_summaries: Dict[str, SingleTaskSummary] = {}
         for task in self.tasks:
-            self.task_summaries[task] = Summary(
-                self.task_loss_fun[task],
-                self.task_metrics[task],
-                self.task_metrics_on_training_set[task],
-                self.task_metrics_on_progress_bar[task],
-                self.monitor,
-                self.mode,
-                task_name=task,
+            self.task_summaries[task] = SingleTaskSummary(
+                metrics = self.task_metrics[task],
+                step_name = step_name,
+                metrics_on_training_set = self.task_metrics_on_training_set[task] if task in self.task_metrics_on_training_set else None,
+                metrics_on_progress_bar = self.task_metrics_on_progress_bar[task] if task in self.task_metrics_on_progress_bar else None,
+                task_name = task,
+                compute_mean = compute_mean,
+                compute_std = compute_std,
             )
 
-        # Current predictor state
-        self.weighted_loss = None
-        self.step_name = None
+    def __getitem__(self, task: str) -> SingleTaskSummary:
+        return self.task_summaries[task]
+    
+    def keys(self) -> List[str]:
+        return self.tasks
 
-    def update_predictor_state(
-        self,
-        step_name: str,
-        targets: Dict[str, Tensor],
-        preds: Dict[str, Tensor],
-        loss: Tensor,
-        task_losses: Dict[str, Tensor],
-        n_epochs: int,
-    ):
+    def update(self, preds: Dict[str, Tensor], targets: Dict[str, Tensor]) -> None:
         r"""
         update the state for all predictors
         Parameters:
-            step_name: the name of the step
             targets: the target tensors
             preds: the prediction tensors
-            loss: the loss tensor
-            task_losses: the task losses
-            n_epochs: the number of epochs
         """
-        self.weighted_loss = loss
-        self.step_name = step_name
         for task in self.tasks:
-            self.task_summaries[task].update_predictor_state(
-                step_name,
-                targets[task],
+            self.task_summaries[task].update(
                 preds[task].detach(),
-                task_losses[task].detach(),
-                n_epochs,
+                targets[task],
             )
 
-    def set_results(self, task_metrics: Dict[str, Dict[str, Tensor]]):
-        """
-        set the results for all tasks
-        Parameters:
-            task_metrics: the metrics for each task
-        """
-        for task in self.tasks:
-            self.task_summaries[task].set_results(task_metrics[task])
-            step_name = self.task_summaries[task].step_name
-            loss = self.task_summaries[task].loss
-            if self.task_summaries[task].is_best_epoch(step_name, loss, task_metrics[task]):
-                self.task_summaries[task].best_summaries[step_name] = self.task_summaries[task].summaries[
-                    step_name
-                ]
-
-    def get_results(
-        self,
-        step_name: str,
-    ) -> Dict[str, Dict[str, Any]]:
-        """
-        retrieve the results
-        Parameters:
-            step_name: the name of the step, i.e. "train"
-        Returns:
-            the results
-        """
-        results = {}
-        for task in self.tasks:
-            results[task] = self.task_summaries[task].get_results(step_name)
-        return results
-
-    def get_best_results(
-        self,
-        step_name: str,
-    ) -> Dict[str, Dict[str, Any]]:
-        """
-        retrieve the best results
-        Parameters:
-            step_name: the name of the step, i.e. "train"
-        Returns:
-            the best results
-        """
-        results = {}
-        for task in self.tasks:
-            results[task] = self.task_summaries[task].get_best_results(step_name)
-        return results
-
     def get_results_on_progress_bar(
         self,
         step_name: str,
@@ -469,78 +365,112 @@ def get_results_on_progress_bar(
         """
         task_results_prog = {}
         for task in self.tasks:
-            # task_results_prog[task] = self.task_summaries[task].get_results_on_progress_bar(step_name)
             task_results_prog.update(self.task_summaries[task].get_results_on_progress_bar(step_name))
         return task_results_prog
 
-    def get_dict_summary(
-        self,
-    ) -> Dict[str, Dict[str, Any]]:
+    def compute(self) -> Dict[str, Tensor]:
         r"""
-        get task summaries in a dictionary
+        compute the metrics for all tasks
         Returns:
-            the task summaries
+            the computed metrics for all tasks
         """
-        task_full_dict = {}
+        computed_metrics = {}
         for task in self.tasks:
-            task_full_dict[task] = self.task_summaries[task].get_dict_summary()
-        return task_full_dict
-
-    def get_metrics_logs(
-        self,
-    ) -> Dict[str, Dict[str, Tensor]]:
+            computed_metrics.update(self.task_summaries[task].compute())
+        return computed_metrics
+    
+    def reset(self) -> None:
         r"""
-        get the logs for the metrics
-        Returns:
-            the task logs for the metrics
+        reset the state of the metrics
         """
-        task_metrics_logs = {}
         for task in self.tasks:
-            task_metrics_logs[task] = self.task_summaries[task].get_metrics_logs()
-            # average metrics
-            for key in task_metrics_logs[task]:
-                if isinstance(task_metrics_logs[task][key], torch.Tensor):
-                    if task_metrics_logs[task][key].numel() > 1:
-                        task_metrics_logs[task][key] = task_metrics_logs[task][key][
-                            task_metrics_logs[task][key] != 0
-                        ].mean()
-
-        # Include global (weighted loss)
-        task_metrics_logs["_global"] = {}
-        task_metrics_logs["_global"][f"loss/{self.step_name}"] = self.weighted_loss.detach().cpu()
-        return task_metrics_logs
-
-    # TODO (Gabriela): This works to fix the logging on TB, but make it more efficient
-    def concatenate_metrics_logs(
-        self,
-        metrics_logs: Dict[str, Dict[str, Tensor]],
-        device: str,
-    ) -> Dict[str, Tensor]:
-        r"""
-        concatenate the metrics logs
-        Parameters:
-            metrics_logs: the metrics logs
-        Returns:
-            the concatenated metrics logs
-        """
-        concatenated_metrics_logs = {}
-        for task in list(self.tasks) + ["_global"]:
-            concatenated_metrics_logs.update(metrics_logs[task])
-        concatenated_metrics_logs[f"loss/{self.step_name}"] = self.weighted_loss.detach().to(device)
-        return concatenated_metrics_logs
+            self.task_summaries[task].reset()
 
-    def metric_log_name(
-        self,
-        task_name: str,
-        metric_name: str,
-        step_name: str,
-    ) -> str:
-        r"""
-        print the metric name, task name and step name
-        Returns:
-            the metric name, task name and step name
-        """
-        if task_name is None:
-            return f"{metric_name}/{step_name}"
+
+class STDMetric(BaseAggregator):
+    """
+    A metric to compute the standard deviation of the predictions or targets.
+    Based on `torchmetrics.Metric`, with a similar implementation to `torchmetric.MeanMetric`.
+
+    Parameters:
+        correction: 
+            The correction to apply to the standard deviation. Instead of dividing by number of samples `N`,
+            we divide by `N-correction`.
+
+        nan_strategy: options:
+            - ``'error'``: if any `nan` values are encountered will give a RuntimeError
+            - ``'warn'``: if any `nan` values are encountered will give a warning and continue
+            - ``'ignore'``: all `nan` values are silently removed
+            - a float: if a float is provided will impute any `nan` values with this value
+
+    """
+    def __init__(self, nan_strategy: Union[Literal["error", "warn", "ignore"], float]="warn", correction:int=0, **kwargs):
+        super().__init__(
+            "sum",
+            default_value=torch.tensor(0.0, dtype=torch.get_default_dtype()),
+            nan_strategy=nan_strategy,
+            state_name="mean_value",
+            **kwargs,
+        )
+        self.add_state("sum", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("sum_of_squares", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total_weight", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.correction = correction
+
+    def update(self, value: Union[float, Tensor], weight: Union[float, Tensor] = 1.0) -> None:
+        if not isinstance(value, Tensor):
+            value = torch.as_tensor(value, dtype=torch.float32, device=value.device)
+        if not isinstance(weight, Tensor):
+            weight = torch.as_tensor(weight, dtype=torch.float32, device=value.device)
+
+        weight = torch.broadcast_to(weight, value.shape).clone()
+        # Check whether `_cast_and_nan_check_input` takes in `weight`
+        if "weight" in inspect.signature(self._cast_and_nan_check_input).parameters:
+            value, weight = self._cast_and_nan_check_input(value, weight)
         else:
-            return f"{task_name}/{metric_name}/{step_name}"
+            weight[value.isnan()] = torch.nan
+            value = self._cast_and_nan_check_input(value)
+            weight = self._cast_and_nan_check_input(weight)
+
+        if value.numel() == 0:
+            return
+
+        self.sum += (value * weight).sum()
+        self.sum_of_squares += (value * value * weight).sum()
+        self.total_weight += weight.sum()
+
+    def compute(self) -> Tensor:
+        dividor = max(0, self.total_weight - self.correction)
+        mean = self.sum / self.total_weight
+        mean_of_squares = self.sum_of_squares / self.total_weight
+        variance = mean_of_squares - mean ** 2
+        variance_corr = variance * (self.total_weight / dividor)
+        return torch.sqrt(variance_corr)
+
+class GradientNormMetric(Metric):
+    """
+    A metric to compute the norm of the gradient.
+    Based on `torchmetrics.Metric`.
+
+    Warning:
+        This metric is not compatible with other metrics since it doesn't take
+        the predictions and targets as input. It takes the model as input.
+        It also doesn't work per task, but for the full model
+    """
+    def __init__(self, dist_sync_on_step=False):
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+        self.add_state("gradient_norm_sq", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total_steps", default=torch.tensor(0.0), dist_reduce_fx="sum")
+
+    def update(self, model: torch.nn.Module) -> None:
+        total_norm = torch.tensor(0.0, device=self.device)
+        for p in model.parameters():
+            if p.grad is not None:
+                param_norm = p.grad.detach().data.norm(2)
+                total_norm += param_norm.detach() ** 2
+        self.gradient_norm_sq += total_norm
+        self.total_steps += 1
+
+    def compute(self) -> Tensor:
+        return (self.gradient_norm_sq / self.total_steps).sqrt()
+
diff --git a/graphium/trainer/progress_bar.py b/graphium/trainer/progress_bar.py
new file mode 100644
index 000000000..a5dde1293
--- /dev/null
+++ b/graphium/trainer/progress_bar.py
@@ -0,0 +1,27 @@
+import sys
+from typing import Any, Callable, Dict, List, Optional, Union, Literal, Iterable
+from lightning.pytorch.callbacks import TQDMProgressBar
+
+
+
+class ProgressBarMetrics(TQDMProgressBar):
+    def __init__(self, metrics_on_progress_bar: Optional[Iterable[str]] = None, loss_alias:Optional[str]="_global/loss/train") -> None:
+        super().__init__()
+        if metrics_on_progress_bar is None:
+            metrics_on_progress_bar = {}
+        self.metrics_on_progress_bar = set(metrics_on_progress_bar)
+        self.loss_alias = loss_alias
+
+    def get_metrics(self, trainer, pl_module) -> Dict[str, Union[int, str, float, Dict[str, float]]]:
+        
+        metrics = super().get_metrics(trainer, pl_module)
+        filtered_metrics = {}
+        for key, metric in metrics.items():
+            if key in self.metrics_on_progress_bar:
+                if key == self.loss_alias:
+                    filtered_metrics["loss"] = metric
+                else:
+                    filtered_metrics[key] = metric
+        
+        return filtered_metrics
+
diff --git a/graphium/utils/spaces.py b/graphium/utils/spaces.py
index 12f07a1ed..c7b6a7ac9 100644
--- a/graphium/utils/spaces.py
+++ b/graphium/utils/spaces.py
@@ -15,7 +15,7 @@
 from copy import deepcopy
 import torch
 import torch.optim.lr_scheduler as sc
-import torchmetrics.functional as TorchMetrics
+import torchmetrics as TorchMetrics
 
 import graphium.nn.base_layers as BaseLayers
 import graphium.nn.ensemble_layers as EnsembleLayers
@@ -23,7 +23,6 @@
 import graphium.utils.custom_lr as CustomLR
 import graphium.data.datamodule as Datamodules
 import graphium.ipu.ipu_losses as IPULosses
-import graphium.ipu.ipu_metrics as Metrics
 import graphium.nn.pyg_layers as PygLayers
 import graphium.nn.residual_connections as Residuals
 import graphium.nn.encoders as Encoders
@@ -102,39 +101,27 @@
 }
 
 METRICS_CLASSIFICATION = {
-    "accuracy": TorchMetrics.accuracy,
-    "averageprecision": TorchMetrics.average_precision,
-    "auroc": TorchMetrics.auroc,
-    "confusionmatrix": TorchMetrics.confusion_matrix,
-    "f1": TorchMetrics.f1_score,
-    "fbeta": TorchMetrics.fbeta_score,
-    "precisionrecallcurve": TorchMetrics.precision_recall_curve,
-    "precision": TorchMetrics.precision,
-    "recall": TorchMetrics.recall,
-    "mcc": TorchMetrics.matthews_corrcoef,
-    "auroc_ipu": Metrics.auroc_ipu,
-    "accuracy_ipu": Metrics.accuracy_ipu,
-    "average_precision_ipu": Metrics.average_precision_ipu,
-    "f1_ipu": Metrics.f1_score_ipu,
-    "fbeta_ipu": Metrics.fbeta_score_ipu,
-    "precision_ipu": Metrics.precision_ipu,
-    "recall_ipu": Metrics.recall_ipu,
+    "accuracy": TorchMetrics.Accuracy,
+    "averageprecision": TorchMetrics.functional.average_precision,  # Not using a class to better handle concatenation of preds and targets
+    "auroc": TorchMetrics.functional.auroc,  # Not using a class to better handle concatenation of preds and targets
+    "confusionmatrix": TorchMetrics.ConfusionMatrix,
+    "f1": TorchMetrics.F1Score,
+    "fbeta": TorchMetrics.FBetaScore,
+    "precisionrecallcurve": TorchMetrics.PrecisionRecallCurve,
+    "precision": TorchMetrics.Precision,
+    "recall": TorchMetrics.Recall,
+    "mcc": TorchMetrics.MatthewsCorrCoef,
 }
 
 METRICS_REGRESSION = {
-    "mae": TorchMetrics.mean_absolute_error,
-    "mape": TorchMetrics.mean_absolute_percentage_error,
-    "mse": TorchMetrics.mean_squared_error,
-    "msle": TorchMetrics.mean_squared_log_error,
-    "pearsonr": TorchMetrics.pearson_corrcoef,
-    "spearmanr": TorchMetrics.spearman_corrcoef,
-    "r2_score": TorchMetrics.r2_score,
-    "cosine": TorchMetrics.cosine_similarity,
-    "pearsonr_ipu": Metrics.pearson_ipu,
-    "spearmanr_ipu": Metrics.spearman_ipu,
-    "r2_score_ipu": Metrics.r2_score_ipu,
-    "mae_ipu": Metrics.mean_absolute_error_ipu,
-    "mse_ipu": Metrics.mean_squared_error_ipu,
+    "mae": TorchMetrics.MeanAbsoluteError,
+    "mape": TorchMetrics.MeanAbsolutePercentageError,
+    "mse": TorchMetrics.MeanSquaredError,
+    "msle": TorchMetrics.MeanSquaredLogError,
+    "pearsonr": TorchMetrics.PearsonCorrCoef,
+    "spearmanr": TorchMetrics.SpearmanCorrCoef,
+    "r2_score": TorchMetrics.R2Score,
+    "cosine": TorchMetrics.CosineSimilarity,
 }
 
 METRICS_DICT = deepcopy(METRICS_CLASSIFICATION)


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 5688c7c94..af85c8616 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -49,6 +49,9 @@ jobs:
       - name: Install library
         run: python -m pip install --no-deps -e . # `-e` required for correct `coverage` run.
 
+      - name: Install test dependencies
+        run: micromamba install -c conda-forge pytdc  # Required to run the `test_finetuning.py`
+
       - name: Install C++ library
         run: cd graphium/graphium_cpp && git clone https://github.com/pybind/pybind11.git && export PYTHONPATH=$PYTHONPATH:./pybind11 && python -m pip install . && cd ../..
 
diff --git a/expts/run_validation_test.py b/expts/run_validation_test.py
index 06804301c..48cb0183e 100644
--- a/expts/run_validation_test.py
+++ b/expts/run_validation_test.py
@@ -8,7 +8,7 @@
 import timeit
 from loguru import logger
 from datetime import datetime
-from pytorch_lightning.utilities.model_summary import ModelSummary
+from lightning.pytorch.utilities.model_summary import ModelSummary
 
 # Current project imports
 import graphium
diff --git a/graphium/cli/train_finetune_test.py b/graphium/cli/train_finetune_test.py
index 09183c69e..4cea1ee8c 100644
--- a/graphium/cli/train_finetune_test.py
+++ b/graphium/cli/train_finetune_test.py
@@ -52,41 +52,6 @@ def cli(cfg: DictConfig) -> None:
     return run_training_finetuning_testing(cfg)
 
 
-def get_replication_factor(cfg):
-    try:
-        ipu_config = cfg.get("accelerator", {}).get("ipu_config", [])
-        for item in ipu_config:
-            if "replicationFactor" in item:
-                # Extract the number between parentheses
-                start = item.find("(") + 1
-                end = item.find(")")
-                if start != 0 and end != -1:
-                    return int(item[start:end])
-    except Exception as e:
-        print(f"An error occurred: {e}")
-
-    # Return default value if replicationFactor is not found or an error occurred
-    return 1
-
-
-def get_gradient_accumulation_factor(cfg):
-    """
-    WARNING: This MUST be called after accelerator overrides have been applied
-    (i.e. after `load_accelerator` has been called)
-    """
-    try:
-        # Navigate through the nested dictionaries and get the gradient accumulation factor
-        grad_accumulation_factor = cfg.get("trainer", {}).get("trainer", {}).get("accumulate_grad_batches", 1)
-
-        # Ensure that the extracted value is an integer
-        return int(grad_accumulation_factor)
-    except Exception as e:
-        print(f"An error occurred: {e}")
-
-    # Return default value if an error occurred
-    return 1
-
-
 def get_training_batch_size(cfg):
     """
     WARNING: This MUST be called after accelerator overrides have been applied
@@ -195,14 +160,6 @@ def run_training_finetuning_testing(cfg: DictConfig) -> None:
         ## Metrics
         metrics = load_metrics(cfg)
 
-        # Note: these MUST be called after `cfg, accelerator = load_accelerator(cfg)`
-        replicas = get_replication_factor(cfg)
-        gradient_acc = get_gradient_accumulation_factor(cfg)
-        micro_bs = get_training_batch_size(cfg)
-        device_iterations = get_training_device_iterations(cfg)
-
-        global_bs = replicas * gradient_acc * micro_bs * device_iterations
-
         ## Predictor
         predictor = load_predictor(
             config=cfg,
@@ -213,17 +170,15 @@ def run_training_finetuning_testing(cfg: DictConfig) -> None:
             accelerator_type=accelerator_type,
             featurization=datamodule.featurization,
             task_norms=datamodule.task_norms,
-            replicas=replicas,
-            gradient_acc=gradient_acc,
-            global_bs=global_bs,
         )
 
     logger.info(predictor.model)
     logger.info(ModelSummary(predictor, max_depth=4))
+    metrics_on_progress_bar = predictor.get_metrics_on_progress_bar
 
     ## Trainer
     date_time_suffix = datetime.now().strftime("%d.%m.%Y_%H.%M.%S")
-    trainer = load_trainer(cfg, accelerator_type, date_time_suffix)
+    trainer = load_trainer(cfg, accelerator_type, date_time_suffix, metrics_on_progress_bar=metrics_on_progress_bar)
 
     if not testing_only:
         # Add the fine-tuning callback to trainer
diff --git a/tests/dummy-pretrained-model.ckpt b/tests/dummy-pretrained-model.ckpt
index b1312cffa..e3df0fa4c 100644
Binary files a/tests/dummy-pretrained-model.ckpt and b/tests/dummy-pretrained-model.ckpt differ
diff --git a/tests/test_data_utils.py b/tests/test_data_utils.py
index 6b73110de..961cef5a7 100644
--- a/tests/test_data_utils.py
+++ b/tests/test_data_utils.py
@@ -18,7 +18,7 @@
 import tempfile
 
 
-class TestDataUtils(ut.TestCase):
+class test_DataUtils(ut.TestCase):
     def test_list_datasets(
         self,
     ):
diff --git a/tests/test_datamodule.py b/tests/test_datamodule.py
index 2cc59b8c4..a2d0c162c 100644
--- a/tests/test_datamodule.py
+++ b/tests/test_datamodule.py
@@ -27,7 +27,7 @@
 TEMP_CACHE_DATA_PATH = "tests/temp_cache_0000"
 
 
-class Test_DataModule(ut.TestCase):
+class test_DataModule(ut.TestCase):
     def test_ogb_datamodule(self):
         # other datasets are too large to be tested
         dataset_names = ["ogbg-molhiv", "ogbg-molpcba", "ogbg-moltox21", "ogbg-molfreesolv"]
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 3624e36ab..feb83a10a 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -105,7 +105,7 @@ def dataframes_to_dataset(dataframes_dict, case_num):
     return multitask_dataset
 
 
-class Test_Multitask_Dataset(ut.TestCase):
+class test_Multitask_Dataset(ut.TestCase):
     # Then we can choose different rows and columns for the tests as we see fit.
     # Remember tests are supposed to be FAST, and reading from the file system multiple times slows things down.
 
diff --git a/tests/test_finetuning.py b/tests/test_finetuning.py
index 52484c4c9..db71ac81b 100644
--- a/tests/test_finetuning.py
+++ b/tests/test_finetuning.py
@@ -16,6 +16,7 @@
 import unittest as ut
 from copy import deepcopy
 from os.path import abspath, dirname
+import shutil
 
 import torch
 from lightning.pytorch.callbacks import Callback
@@ -29,7 +30,6 @@
     load_metrics,
     load_predictor,
     load_trainer,
-    save_params_to_wandb,
 )
 from graphium.finetuning import GraphFinetuning, modify_cfg_for_finetuning
 from graphium.trainer import PredictorModule
@@ -40,7 +40,7 @@
 os.chdir(MAIN_DIR)
 
 
-class Test_Finetuning(ut.TestCase):
+class test_Finetuning(ut.TestCase):
     def test_finetuning_from_task_head(self):
         # Skip test if PyTDC package not installed
         try:
@@ -60,9 +60,14 @@ def test_finetuning_from_task_head(self):
         # Initialize the accelerator
         cfg, accelerator_type = load_accelerator(cfg)
 
+        # If the data_cache directory exists, delete it for the purpose of the test
+        data_cache = cfg["datamodule"]["args"]["processed_graph_data_path"]
+        if os.path.exists(data_cache):
+            shutil.rmtree(data_cache)
+
         # Load and initialize the dataset
         datamodule = load_datamodule(cfg, accelerator_type)
-        datamodule.task_specific_args["lipophilicity_astrazeneca"].sample_size = 100
+        datamodule.task_specific_args["lipophilicity_astrazeneca"].sample_size = 300
 
         # Initialize the network
         model_class, model_kwargs = load_architecture(
@@ -149,7 +154,7 @@ def test_finetuning_from_task_head(self):
         #################################################
 
         # Define test callback that checks for correct (un)freezing
-        class TestCallback(Callback):
+        class CallbackTesting(Callback):
             def __init__(self, cfg):
                 super().__init__()
 
@@ -217,18 +222,20 @@ def on_train_epoch_start(self, trainer, pl_module):
 
                     assert not False in unfrozen_parameters
 
-        trainer = load_trainer(cfg, accelerator_type)
+        metrics_on_progress_bar = predictor.get_metrics_on_progress_bar
+        trainer = load_trainer(cfg, accelerator_type, metrics_on_progress_bar=metrics_on_progress_bar)
 
         finetuning_training_kwargs = cfg["finetuning"]["training_kwargs"]
         trainer.callbacks.append(GraphFinetuning(**finetuning_training_kwargs))
 
         # Add test callback to trainer
-        trainer.callbacks.append(TestCallback(cfg))
+        trainer.callbacks.append(CallbackTesting(cfg))
 
         predictor.set_max_nodes_edges_per_graph(datamodule, stages=["train", "val"])
 
         # Run the model training
         trainer.fit(model=predictor, datamodule=datamodule)
+        trainer.test(model=predictor, datamodule=datamodule)
 
     def test_finetuning_from_gnn(self):
         # Skip test if PyTDC package not installed
@@ -249,9 +256,14 @@ def test_finetuning_from_gnn(self):
         # Initialize the accelerator
         cfg, accelerator_type = load_accelerator(cfg)
 
+        # If the data_cache directory exists, delete it for the purpose of the test
+        data_cache = cfg["datamodule"]["args"]["processed_graph_data_path"]
+        if os.path.exists(data_cache):
+            shutil.rmtree(data_cache)
+
         # Load and initialize the dataset
         datamodule = load_datamodule(cfg, accelerator_type)
-        datamodule.task_specific_args["lipophilicity_astrazeneca"].sample_size = 100
+        datamodule.task_specific_args["lipophilicity_astrazeneca"].sample_size = 300
 
         # Initialize the network
         model_class, model_kwargs = load_architecture(
@@ -335,7 +347,7 @@ def test_finetuning_from_gnn(self):
         #################################################
 
         # Define test callback that checks for correct (un)freezing
-        class TestCallback(Callback):
+        class CallbackTesting(Callback):
             def __init__(self, cfg):
                 super().__init__()
 
@@ -392,18 +404,20 @@ def on_train_epoch_start(self, trainer, pl_module):
 
                     assert not False in unfrozen_parameters
 
-        trainer = load_trainer(cfg, accelerator_type)
+        metrics_on_progress_bar = predictor.get_metrics_on_progress_bar
+        trainer = load_trainer(cfg, accelerator_type, metrics_on_progress_bar=metrics_on_progress_bar)
 
         finetuning_training_kwargs = cfg["finetuning"]["training_kwargs"]
         trainer.callbacks.append(GraphFinetuning(**finetuning_training_kwargs))
 
         # Add test callback to trainer
-        trainer.callbacks.append(TestCallback(cfg))
+        trainer.callbacks.append(CallbackTesting(cfg))
 
         predictor.set_max_nodes_edges_per_graph(datamodule, stages=["train", "val"])
 
         # Run the model training
         trainer.fit(model=predictor, datamodule=datamodule)
+        trainer.test(model=predictor, datamodule=datamodule)
 
 
 if __name__ == "__main__":
diff --git a/tests/test_ipu_dataloader.py b/tests/test_ipu_dataloader.py
index 39f754e20..2b018206e 100644
--- a/tests/test_ipu_dataloader.py
+++ b/tests/test_ipu_dataloader.py
@@ -47,7 +47,7 @@ def global_batch_collator(batch_size, batches):
 
 @pytest.mark.ipu
 class test_DataLoading(ut.TestCase):
-    class TestSimpleLightning(LightningModule):
+    class SimpleLightningTesting(LightningModule):
         # Create a basic Ligthning for testing the batch sizes
         def __init__(self, batch_size, node_feat_size, edge_feat_size, num_batch) -> None:
             super().__init__()
@@ -98,7 +98,7 @@ def assert_shapes(self, batch, batch_idx, step):
         def configure_optimizers(self):
             return torch.optim.Adam(self.parameters(), lr=1e-3)
 
-    class TestDataset(torch.utils.data.Dataset):
+    class DatasetTesting(torch.utils.data.Dataset):
         # Create a simple dataset for testing the Lightning integration
         def __init__(self, labels, node_features, edge_features):
             self.labels = labels
diff --git a/tests/test_ipu_metrics.py b/tests/test_ipu_metrics.py
deleted file mode 100644
index ee4801e7b..000000000
--- a/tests/test_ipu_metrics.py
+++ /dev/null
@@ -1,774 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-import unittest as ut
-import torch
-from torchmetrics.functional import (
-    auroc,
-    average_precision,
-    precision,
-    accuracy,
-    recall,
-    pearson_corrcoef,
-    spearman_corrcoef,
-    r2_score,
-    f1_score,
-    fbeta_score,
-    mean_squared_error,
-    mean_absolute_error,
-)
-from copy import deepcopy
-import pytest
-
-from graphium.ipu.ipu_metrics import (
-    auroc_ipu,
-    average_precision_ipu,
-    precision_ipu,
-    accuracy_ipu,
-    recall_ipu,
-    pearson_ipu,
-    spearman_ipu,
-    r2_score_ipu,
-    f1_score_ipu,
-    fbeta_score_ipu,
-    mean_squared_error_ipu,
-    mean_absolute_error_ipu,
-)
-
-
-@pytest.mark.ipu
-class test_Metrics(ut.TestCase):
-    torch.manual_seed(42)
-    preds = torch.rand((100, 10), dtype=torch.float32)
-    target = torch.rand((100, 10), dtype=torch.float32)
-
-    th = 0.7
-    nan_th = 0.2
-    preds_greater = preds > th
-    target_greater = (target > th).to(torch.float32)
-    target_greater_nan = deepcopy(target_greater)
-    is_nan = target < nan_th
-    target_greater_nan[target < nan_th] = torch.nan
-    target_nan = deepcopy(target)
-    target_nan[target < nan_th] = torch.nan
-
-    def test_auroc(self):
-        preds = deepcopy(self.preds)[:, 0]
-        target = deepcopy(self.target)[:, 0]
-        target_nan = deepcopy(self.target_nan)[:, 0]
-
-        target[target < 0.5] = 0
-        target[target >= 0.5] = 1
-
-        target_nan[target_nan < 0.5] = 0
-        target_nan[target_nan >= 0.5] = 1
-
-        # Regular loss
-        score_true = auroc(preds, target.to(int))
-        score_ipu = auroc_ipu(preds, target)
-        self.assertFalse(score_true.isnan(), "Regular AUROC score is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Regular AUROC score is different"
-        )
-
-        # Weighted loss (As in BCE)
-        sample_weights = torch.rand(preds.shape[0], dtype=torch.float32)
-        score_true = auroc(preds, target.to(int), sample_weights=sample_weights)
-        score_ipu = auroc_ipu(preds, target, sample_weights=sample_weights)
-        self.assertFalse(score_true.isnan(), "Regular AUROC score is NaN")
-        self.assertAlmostEqual(score_true.item(), score_ipu.item(), msg="Weighted AUROC score is different")
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = auroc(preds[not_nan], target[not_nan].to(int))
-        score_ipu = auroc_ipu(preds, target_nan)
-        self.assertFalse(score_true.isnan(), "Regular AUROC score with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Regular AUROCIPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Regular AUROC score with NaN is different"
-        )
-
-        # Weighted loss with NaNs in target (As in BCE)
-        not_nan = ~target_nan.isnan()
-        sample_weights = torch.rand(preds.shape, dtype=torch.float32)
-        loss_true = auroc(preds[not_nan], target_nan[not_nan].to(int), sample_weights=sample_weights[not_nan])
-        loss_ipu = auroc_ipu(preds, target_nan, sample_weights=sample_weights)
-        self.assertFalse(loss_true.isnan(), "Weighted AUROC score with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Weighted AUROC IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            # AssertionError: 0.6603766679763794 != 0.6234951615333557 within 2 places
-            loss_true.item(),
-            loss_ipu.item(),
-            places=6,
-            msg="Weighted AUROC with NaN is different",
-        )
-
-    def test_average_precision(self):  # TODO: Make work with multi-class
-        preds = deepcopy(self.preds)[:, 0]
-        target = deepcopy(self.target)[:, 0]
-        target_nan = deepcopy(self.target_nan)[:, 0]
-
-        target[target < 0.5] = 0
-        target[target >= 0.5] = 1
-
-        target_nan[target_nan < 0.5] = 0
-        target_nan[target_nan >= 0.5] = 1
-
-        # Regular loss
-        score_true = average_precision(preds, target.to(int), task="binary")
-        score_ipu = average_precision_ipu(preds, target.to(int), task="binary")
-        self.assertFalse(score_true.isnan(), "Regular Average Precision is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Regular Average Precision is different"
-        )
-
-        # Regular average precision with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = average_precision(preds[not_nan], target[not_nan].to(int), task="binary")
-        score_ipu = average_precision_ipu(preds, target_nan, task="binary")
-        self.assertFalse(score_true.isnan(), "Regular Average Precision with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Regular Average Precision IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(),
-            score_ipu.item(),
-            places=6,
-            msg="Regular Average Precision with NaN is different",
-        )
-
-    def test_precision(self):
-        preds = deepcopy(self.preds)[:, :4]
-        target = deepcopy(self.target)[:, 0]
-        t = deepcopy(target)
-
-        target[t < 0.4] = 0
-        target[(t >= 0.4) & (t < 0.6)] = 1
-        target[(t >= 0.6) & (t < 0.8)] = 2
-        target[(t >= 0.8)] = 3
-
-        target_nan = deepcopy(target)
-        target_nan[self.is_nan[:, 0]] = float("nan")
-        target_nan_bin = deepcopy(target_nan)
-        target_nan_bin[target_nan > 0] = 1
-
-        # Micro precision binary
-        score_true = precision(preds[:, 0], target.to(int) > 0, average="micro")
-        score_ipu = precision_ipu(preds[:, 0], target > 0, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Precision binary is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision binary is different"
-        )
-
-        # Micro precision binary with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = precision(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
-        score_ipu = precision_ipu(preds[:, 0], target_nan_bin, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Precision binary with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Micro Precision binary IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision with NaN is different"
-        )
-
-        # Micro precision
-        score_true = precision(preds, target.to(int), average="micro")
-        score_ipu = precision_ipu(preds, target, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Precision is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision is different"
-        )
-
-        # Micro precision with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = precision(preds[not_nan], target[not_nan].to(int), average="micro")
-        score_ipu = precision_ipu(preds, target_nan, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Precision with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Micro Precision IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Precision with NaN is different"
-        )
-
-        # Macro precision
-        score_true = precision(preds, target.to(int), average="macro", num_classes=4)
-        score_ipu = precision_ipu(preds, target, average="macro", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Macro Precision is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Macro Precision is different"
-        )
-
-        # Macro precision multiclass with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = precision(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
-        score_ipu = precision_ipu(preds, target_nan, average="macro", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Macro Precision multiclass with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Macro Precision multiclass IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(),
-            score_ipu.item(),
-            places=6,
-            msg="Macro Precision multiclass with NaN is different",
-        )
-
-        # Macro precision multiclass with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = precision(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
-        score_ipu = precision_ipu(preds, target_nan, average="macro", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Macro Precision multiclass with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Macro Precision multiclass IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(),
-            score_ipu.item(),
-            places=6,
-            msg="Macro Precision multiclass with NaN is different",
-        )
-
-        # Weighted precision multiclass
-        score_true = precision(preds, target.to(int), average="weighted", num_classes=4)
-        score_ipu = precision_ipu(preds, target, average="weighted", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Weighted Precision multiclass is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Weighted Precision multiclass is different"
-        )
-
-        # Weighted precision multiclass with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = precision(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
-        score_ipu = precision_ipu(preds, target_nan, average="weighted", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Weighted Precision multiclass with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Weighted Precision multiclass IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(),
-            score_ipu.item(),
-            places=6,
-            msg="Regular Average Precision multiclass with NaN is different",
-        )
-
-    def test_accuracy(self):
-        preds = deepcopy(self.preds)[:, :4]
-        target = deepcopy(self.target)[:, 0]
-        t = deepcopy(target)
-
-        target[t < 0.4] = 0
-        target[(t >= 0.4) & (t < 0.6)] = 1
-        target[(t >= 0.6) & (t < 0.8)] = 2
-        target[(t >= 0.8)] = 3
-
-        target_nan = deepcopy(target)
-        target_nan[self.is_nan[:, 0]] = float("nan")
-        target_nan_bin = deepcopy(target_nan)
-        target_nan_bin[target_nan > 0] = 1
-
-        # Micro accuracy binary
-        score_true = accuracy(preds[:, 0], target.to(int) > 0, average="micro")
-        score_ipu = accuracy_ipu(preds[:, 0], target > 0, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Accuracy binary is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy binary is different"
-        )
-
-        # Micro accuracy binary with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = accuracy(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
-        score_ipu = accuracy_ipu(preds[:, 0], target_nan_bin, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Accuracy binary with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Micro Accuracy binary IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy with NaN is different"
-        )
-
-        # Micro accuracy
-        score_true = accuracy(preds, target.to(int), average="micro")
-        score_ipu = accuracy_ipu(preds, target, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Accuracy is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy is different"
-        )
-
-        # Micro accuracy with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = accuracy(preds[not_nan], target[not_nan].to(int), average="micro")
-        score_ipu = accuracy_ipu(preds, target_nan, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Accuracy with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Micro Accuracy IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Accuracy with NaN is different"
-        )
-
-        # Macro accuracy
-        score_true = accuracy(preds, target.to(int), average="macro", num_classes=4)
-        score_ipu = accuracy_ipu(preds, target, average="macro", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Macro Accuracy is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Macro Accuracy is different"
-        )
-
-        # Macro accuracy with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = accuracy(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
-        score_ipu = accuracy_ipu(preds, target_nan, average="macro", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Macro Accuracy with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Macro Accuracy IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Macro Accuracy with NaN is different"
-        )
-
-        # Weighted accuracy
-        score_true = accuracy(preds, target.to(int), average="weighted", num_classes=4)
-        score_ipu = accuracy_ipu(preds, target, average="weighted", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Weighted Accuracy is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Weighted Accuracy is different"
-        )
-
-        # Weighted accuracy with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = accuracy(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
-        score_ipu = accuracy_ipu(preds, target_nan, average="weighted", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Weighted Accuracy with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Weighted Accuracy IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Regular Accuracy with NaN is different"
-        )
-
-    def test_recall(self):
-        preds = deepcopy(self.preds)[:, :4]
-        target = deepcopy(self.target)[:, 0]
-        t = deepcopy(target)
-
-        target[t < 0.4] = 0
-        target[(t >= 0.4) & (t < 0.6)] = 1
-        target[(t >= 0.6) & (t < 0.8)] = 2
-        target[(t >= 0.8)] = 3
-
-        target_nan = deepcopy(target)
-        target_nan[self.is_nan[:, 0]] = float("nan")
-        target_nan_bin = deepcopy(target_nan)
-        target_nan_bin[target_nan > 0] = 1
-
-        # Micro recall binary
-        score_true = recall(preds[:, 0], target.to(int) > 0, average="micro")
-        score_ipu = recall_ipu(preds[:, 0], target > 0, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Recall binary is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Recall binary is different"
-        )
-
-        # Micro recall binary with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = recall(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
-        score_ipu = recall_ipu(preds[:, 0], target_nan_bin, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Recall binary with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Micro Recall binary IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Recall binary with NaN is different"
-        )
-
-        # Micro recall
-        score_true = recall(preds, target.to(int), average="micro")
-        score_ipu = recall_ipu(preds, target, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Recall is NaN")
-        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=6, msg="Micro Recall is different")
-
-        # Micro recall with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = recall(preds[not_nan], target[not_nan].to(int), average="micro")
-        score_ipu = recall_ipu(preds, target_nan, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro Recall with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Micro Recall IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro Recall with NaN is different"
-        )
-
-        # Macro recall multiclass
-        score_true = recall(preds, target.to(int), average="macro", num_classes=4)
-        score_ipu = recall_ipu(preds, target, average="macro", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Macro Recall is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Macro Recall multiclass is different"
-        )
-
-        # Macro recall multiclass with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = recall(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
-        score_ipu = recall_ipu(preds, target_nan, average="macro", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Macro Recall multiclass with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Macro Recall multiclass IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Macro Recall multiclass with NaN is different"
-        )
-
-        # Weighted recallmulticlass
-        score_true = recall(preds, target.to(int), average="weighted", num_classes=4)
-        score_ipu = recall_ipu(preds, target, average="weighted", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Weighted Recall is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Weighted Recall is different"
-        )
-
-        # Weighted recall multiclass with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = recall(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
-        score_ipu = recall_ipu(preds, target_nan, average="weighted", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Weighted Recall multiclass with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Weighted Recall multiclass IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(),
-            score_ipu.item(),
-            places=6,
-            msg="Regular Recall multiclass with NaN is different",
-        )
-
-    def test_pearsonr(self):
-        preds = deepcopy(self.preds)[:, 0]
-        target = deepcopy(self.target)[:, 0] + preds
-        target_nan = deepcopy(target)
-        target_nan[self.is_nan[:, 0]] = float("nan")
-
-        # Regular loss
-        score_true = pearson_corrcoef(preds, target)
-        score_ipu = pearson_ipu(preds, target)
-        self.assertFalse(score_true.isnan(), "Pearson is NaN")
-        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=4, msg="Pearson is different")
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = pearson_corrcoef(preds[not_nan], target[not_nan])
-        score_ipu = pearson_ipu(preds, target_nan)
-        self.assertFalse(score_true.isnan(), "Regular PearsonR with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "IPU PearsonR score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=4, msg="Pearson with NaN is different"
-        )
-
-    def test_spearmanr(self):
-        preds = deepcopy(self.preds)[:, 0]
-        target = deepcopy(self.target)[:, 0] + preds
-        target_nan = deepcopy(target)
-        target_nan[self.is_nan[:, 0]] = float("nan")
-
-        # Regular loss
-        score_true = spearman_corrcoef(preds, target)
-        score_ipu = spearman_ipu(preds, target)
-        self.assertFalse(score_true.isnan(), "Spearman is NaN")
-        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=4, msg="Spearman is different")
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = spearman_corrcoef(preds[not_nan], target[not_nan])
-        score_ipu = spearman_ipu(preds, target_nan)
-        self.assertFalse(score_true.isnan(), "Regular Spearman with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "IPU Spearman score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=4, msg="Spearman with NaN is different"
-        )
-
-    def test_r2_score(self):
-        preds = deepcopy(self.preds)
-        target = deepcopy(self.target) + preds
-        target_nan = deepcopy(target)
-        target_nan[self.is_nan] = float("nan")
-
-        # Regular loss
-        score_true = r2_score(preds, target)
-        score_ipu = r2_score_ipu(preds, target)
-        self.assertFalse(score_true.isnan(), "r2_score is NaN")
-        self.assertAlmostEqual(score_true.item(), score_ipu.item(), places=4, msg="r2_score is different")
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_ipu = r2_score_ipu(preds, target_nan, multioutput="raw_values")
-        for ii in range(preds.shape[1]):
-            score_true = r2_score(
-                preds[:, ii][not_nan[:, ii]], target_nan[:, ii][not_nan[:, ii]], multioutput="raw_values"
-            )
-            self.assertFalse(score_true.isnan().any(), f"{ii}: r2_score with target_nan is NaN")
-            self.assertFalse(score_ipu[ii].isnan().any(), f"{ii}: IPU r2_score with target_nan is NaN")
-            self.assertAlmostEqual(
-                score_true.item(), score_ipu[ii].item(), places=4, msg=f"{ii}: r2_score with NaN is different"
-            )
-
-    def test_fbeta_score(self):
-        preds = deepcopy(self.preds)[:, :4]
-        target = deepcopy(self.target)[:, 0]
-        t = deepcopy(target)
-
-        target[t < 0.4] = 0
-        target[(t >= 0.4) & (t < 0.6)] = 1
-        target[(t >= 0.6) & (t < 0.8)] = 2
-        target[(t >= 0.8)] = 3
-
-        target_nan = deepcopy(target)
-        target_nan[self.is_nan[:, 0]] = float("nan")
-        target_nan_bin = deepcopy(target_nan)
-        target_nan_bin[target_nan > 0] = 1
-
-        # Micro fbeta_score binary
-        score_true = fbeta_score(preds[:, 0], target.to(int) > 0, average="micro", beta=0.5)
-        score_ipu = fbeta_score_ipu(preds[:, 0], target > 0, average="micro", beta=0.5)
-        self.assertFalse(score_true.isnan(), "Micro FBETA_score binary is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro FBETA_score binary is different"
-        )
-
-        # Micro fbeta_score binary with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = fbeta_score(
-            preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro", beta=0.5
-        )
-        score_ipu = fbeta_score_ipu(preds[:, 0], target_nan_bin, average="micro", beta=0.5)
-        self.assertFalse(score_true.isnan(), "Micro FBETA_score binary with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Micro FBETA_score binary IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(),
-            score_ipu.item(),
-            places=6,
-            msg="Micro FBETA_score binary with NaN is different",
-        )
-
-        # Micro fbeta_score
-        score_true = fbeta_score(preds, target.to(int), average="micro", beta=0.5)
-        score_ipu = fbeta_score_ipu(preds, target, average="micro", beta=0.5)
-        self.assertFalse(score_true.isnan(), "Micro FBETA_score is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro FBETA_score is different"
-        )
-
-        # Micro fbeta_score with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = fbeta_score(preds[not_nan], target[not_nan].to(int), average="micro", beta=0.5)
-        score_ipu = fbeta_score_ipu(preds, target_nan, average="micro", beta=0.5)
-        self.assertFalse(score_true.isnan(), "Micro FBETA_score with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Micro FBETA_score IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro FBETA_score with NaN is different"
-        )
-
-        # Macro fbeta_score multiclass
-        score_true = fbeta_score(preds, target.to(int), average="macro", num_classes=4, beta=0.5)
-        score_ipu = fbeta_score_ipu(preds, target, average="macro", num_classes=4, beta=0.5)
-        self.assertFalse(score_true.isnan(), "Macro FBETA_score is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Macro FBETA_score multiclass is different"
-        )
-
-        # Macro fbeta_score multiclass with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = fbeta_score(
-            preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4, beta=0.5
-        )
-        score_ipu = fbeta_score_ipu(preds, target_nan, average="macro", num_classes=4, beta=0.5)
-        self.assertFalse(score_true.isnan(), "Macro FBETA_score multiclass with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Macro FBETA_score multiclass IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(),
-            score_ipu.item(),
-            places=6,
-            msg="Macro FBETA_score multiclass with NaN is different",
-        )
-
-        # Weighted fbeta_scoremulticlass
-        score_true = fbeta_score(preds, target.to(int), average="weighted", num_classes=4, beta=0.5)
-        score_ipu = fbeta_score_ipu(preds, target, average="weighted", num_classes=4, beta=0.5)
-        self.assertFalse(score_true.isnan(), "Weighted FBETA_score is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Weighted FBETA_score is different"
-        )
-
-        # Weighted fbeta_score multiclass with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = fbeta_score(
-            preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4, beta=0.5
-        )
-        score_ipu = fbeta_score_ipu(preds, target_nan, average="weighted", num_classes=4, beta=0.5)
-        self.assertFalse(score_true.isnan(), "Weighted FBETA_score multiclass with target_nan is NaN")
-        self.assertFalse(
-            score_ipu.isnan(), "Weighted FBETA_score multiclass IPU score with target_nan is NaN"
-        )
-        self.assertAlmostEqual(
-            score_true.item(),
-            score_ipu.item(),
-            places=6,
-            msg="Regular FBETA_score multiclass with NaN is different",
-        )
-
-    def test_f1_score(self):
-        preds = deepcopy(self.preds)[:, :4]
-        target = deepcopy(self.target)[:, 0]
-        t = deepcopy(target)
-
-        target[t < 0.4] = 0
-        target[(t >= 0.4) & (t < 0.6)] = 1
-        target[(t >= 0.6) & (t < 0.8)] = 2
-        target[(t >= 0.8)] = 3
-
-        target_nan = deepcopy(target)
-        target_nan[self.is_nan[:, 0]] = float("nan")
-        target_nan_bin = deepcopy(target_nan)
-        target_nan_bin[target_nan > 0] = 1
-
-        # Micro f1_score binary
-        score_true = f1_score(preds[:, 0], target.to(int) > 0, average="micro")
-        score_ipu = f1_score_ipu(preds[:, 0], target > 0, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro F1_score binary is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score binary is different"
-        )
-
-        # Micro f1_score binary with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = f1_score(preds[:, 0][not_nan], target_nan_bin[not_nan].to(int), average="micro")
-        score_ipu = f1_score_ipu(preds[:, 0], target_nan_bin, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro F1_score binary with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Micro F1_score binary IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score binary with NaN is different"
-        )
-
-        # Micro f1_score
-        score_true = f1_score(preds, target.to(int), average="micro")
-        score_ipu = f1_score_ipu(preds, target, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro F1_score is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score is different"
-        )
-
-        # Micro f1_score with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = f1_score(preds[not_nan], target[not_nan].to(int), average="micro")
-        score_ipu = f1_score_ipu(preds, target_nan, average="micro")
-        self.assertFalse(score_true.isnan(), "Micro F1_score with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Micro F1_score IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Micro F1_score with NaN is different"
-        )
-
-        # Macro f1_score multiclass
-        score_true = f1_score(preds, target.to(int), average="macro", num_classes=4)
-        score_ipu = f1_score_ipu(preds, target, average="macro", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Macro F1_score is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Macro F1_score multiclass is different"
-        )
-
-        # Macro f1_score multiclass with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = f1_score(preds[not_nan], target[not_nan].to(int), average="macro", num_classes=4)
-        score_ipu = f1_score_ipu(preds, target_nan, average="macro", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Macro F1_score multiclass with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Macro F1_score multiclass IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(),
-            score_ipu.item(),
-            places=6,
-            msg="Macro F1_score multiclass with NaN is different",
-        )
-
-        # Weighted f1_scoremulticlass
-        score_true = f1_score(preds, target.to(int), average="weighted", num_classes=4)
-        score_ipu = f1_score_ipu(preds, target, average="weighted", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Weighted F1_score is NaN")
-        self.assertAlmostEqual(
-            score_true.item(), score_ipu.item(), places=6, msg="Weighted F1_score is different"
-        )
-
-        # Weighted f1_score multiclass with NaNs in target
-        not_nan = ~target_nan.isnan()
-        score_true = f1_score(preds[not_nan], target[not_nan].to(int), average="weighted", num_classes=4)
-        score_ipu = f1_score_ipu(preds, target_nan, average="weighted", num_classes=4)
-        self.assertFalse(score_true.isnan(), "Weighted F1_score multiclass with target_nan is NaN")
-        self.assertFalse(score_ipu.isnan(), "Weighted F1_score multiclass IPU score with target_nan is NaN")
-        self.assertAlmostEqual(
-            score_true.item(),
-            score_ipu.item(),
-            places=6,
-            msg="Regular F1_score multiclass with NaN is different",
-        )
-
-    def test_mse(self):
-        preds = deepcopy(self.preds)
-        target = deepcopy(self.target)
-        target_nan = deepcopy(self.target_nan)
-        squared = True
-
-        # Regular loss
-        loss_true = mean_squared_error(preds, target, squared)
-        loss_ipu = mean_squared_error_ipu(preds=preds, target=target, squared=squared)
-        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular Mean Squared Error is different"
-        )
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        loss_true = mean_squared_error(preds[not_nan], target[not_nan], squared)
-        loss_ipu = mean_squared_error_ipu(preds=preds, target=target_nan, squared=squared)
-        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Regular Mean Squared Error IPU with target_nan is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(),
-            loss_ipu.item(),
-            places=6,
-            msg="Regular Mean Squared Error with NaN is different",
-        )
-
-        squared = False
-
-        # Regular loss
-        loss_true = mean_squared_error(preds, target, squared)
-        loss_ipu = mean_squared_error_ipu(preds=preds, target=target, squared=squared)
-        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular Mean Squared Error is different"
-        )
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        loss_true = mean_squared_error(preds[not_nan], target[not_nan], squared)
-        loss_ipu = mean_squared_error_ipu(preds=preds, target=target_nan, squared=squared)
-        self.assertFalse(loss_true.isnan(), "Regular Mean Squared Error with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Regular Mean Squared Error IPU with target_nan is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(),
-            loss_ipu.item(),
-            places=6,
-            msg="Regular Mean Squared Error with NaN is different",
-        )
-
-    def test_mae(self):
-        preds = deepcopy(self.preds)
-        target = deepcopy(self.target)
-        target_nan = deepcopy(self.target_nan)
-
-        # Regular loss
-        loss_true = mean_absolute_error(preds, target)
-        loss_ipu = mean_absolute_error_ipu(preds=preds, target=target)
-        self.assertFalse(loss_true.isnan(), "Regular Mean Absolute Error is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular Mean Absolute Error is different"
-        )
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        loss_true = mean_absolute_error(preds[not_nan], target[not_nan])
-        loss_ipu = mean_absolute_error_ipu(preds=preds, target=target_nan)
-        self.assertFalse(loss_true.isnan(), "Regular Mean Absolute Error with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Regular Mean Absolute Error IPU with target_nan is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(),
-            loss_ipu.item(),
-            places=6,
-            msg="Regular Mean Absolute Error with NaN is different",
-        )
-
-
-if __name__ == "__main__":
-    ut.main()
diff --git a/tests/test_ipu_to_dense_batch.py b/tests/test_ipu_to_dense_batch.py
index f2e29537b..157b3b647 100644
--- a/tests/test_ipu_to_dense_batch.py
+++ b/tests/test_ipu_to_dense_batch.py
@@ -24,7 +24,7 @@
 
 
 @pytest.mark.ipu
-class TestIPUBatch:
+class test_IPUBatch:
     @pytest.fixture(autouse=True)
     def setup_class(self):
         self.in_dim = 12
diff --git a/tests/test_loaders.py b/tests/test_loaders.py
index 22611f32f..d4bfbc993 100644
--- a/tests/test_loaders.py
+++ b/tests/test_loaders.py
@@ -17,7 +17,7 @@
 import unittest as ut
 
 
-class TestLoader(ut.TestCase):
+class test_Loader(ut.TestCase):
     def test_merge_dicts(self):
         dict_a = {"a": {"b": {"c": 1, "d": 2}, "e": 3}, "f": 4}
 
diff --git a/tests/test_losses.py b/tests/test_losses.py
index b2f343bf9..82e8090ea 100644
--- a/tests/test_losses.py
+++ b/tests/test_losses.py
@@ -26,7 +26,8 @@
 
 def _parse(loss_fun):
     eval_options = EvalOptions(loss_fun=loss_fun, metrics_on_progress_bar=None)
-    return eval_options.parse_loss_fun(loss_fun)
+    loss_name, loss_fun = eval_options.parse_loss_fun(loss_fun)
+    return loss_fun
 
 
 class test_HybridCELoss(ut.TestCase):
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index dc5bc01b2..cdeac5f7d 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -27,7 +27,8 @@
     Thresholder,
 )
 
-from torchmetrics.functional import mean_squared_error
+from torchmetrics.functional import mean_squared_error, pearson_corrcoef, auroc
+from torchmetrics import MeanSquaredError
 
 
 class test_Metrics(ut.TestCase):
@@ -142,12 +143,10 @@ def test_target_nan_mask(self):
 
     def test_pickling(self):
         pickle_file = os.path.join(tempfile.gettempdir(), "test_metric_pickled.pkl")
-        metrics = ["mae", "mse", mean_squared_error]
+        metrics = ["mae", "mse", MeanSquaredError]
         target_nan_masks = [None, 2, "ignore"]
         multitask_handlings = [None, "flatten", "mean-per-label"]
-        squeeze_targets = [True, False]
         target_to_ints = [True, False]
-        other_kwargs = [{}, {"squared": False}]
         thresholds = [
             None,
             {"threshold": 0.2, "operator": "greater"},
@@ -159,73 +158,134 @@ def test_pickling(self):
 
         for metric in metrics:
             for target_nan_mask in target_nan_masks:
-                for kwargs in other_kwargs:
-                    for threshold_kwargs in thresholds:
-                        for multitask_handling in multitask_handlings:
-                            for squeeze_target in squeeze_targets:
-                                for target_to_int in target_to_ints:
-                                    err_msg = f"{metric} - {target_nan_mask} - {kwargs} - {threshold_kwargs}"
-
-                                    if (multitask_handling is None) and (target_nan_mask == "ignore"):
-                                        # Raise with incompatible options
-                                        with self.assertRaises(ValueError):
-                                            MetricWrapper(
-                                                metric=metric,
-                                                threshold_kwargs=threshold_kwargs,
-                                                target_nan_mask=target_nan_mask,
-                                                multitask_handling=multitask_handling,
-                                                squeeze_target=squeeze_target,
-                                                target_to_int=target_to_int,
-                                                **kwargs,
-                                            )
-
-                                    else:
-                                        metric_wrapper = MetricWrapper(
-                                            metric=metric,
-                                            threshold_kwargs=threshold_kwargs,
-                                            target_nan_mask=target_nan_mask,
-                                            multitask_handling=multitask_handling,
-                                            squeeze_target=squeeze_target,
-                                            target_to_int=target_to_int,
-                                            **kwargs,
-                                        )
-
-                                        # Check that the metric can be saved and re-loaded without error
-                                        torch.save(metric_wrapper, pickle_file)
-                                        metric_wrapper2 = torch.load(pickle_file)
-                                        self.assertTrue(metric_wrapper == metric_wrapper2, msg=err_msg)
-
-                                        # Check that the metric only contains primitive types
-                                        state = metric_wrapper.__getstate__()
-                                        if state["threshold_kwargs"] is not None:
-                                            self.assertIsInstance(
-                                                state["threshold_kwargs"], dict, msg=err_msg
-                                            )
-                                        if isinstance(metric, str):
-                                            self.assertIsInstance(state["metric"], str, msg=err_msg)
-
-    def test_classifigression_target_squeezing(self):
-        preds = torch.Tensor([[0.1, 0.1, 0.3, 0.5, 0.0, 0.1, 0.0, 0.7, 0.2, 0.0]])
-        target = torch.Tensor([3, 0])
-        expected_scores = [0.5, 0.75]
-        n_brackets = 5
-        metrics = ["accuracy", "averageprecision"]
-        other_kwargs = [
-            {"task": "multiclass", "num_classes": n_brackets, "top_k": 1},
-            {"task": "multiclass", "num_classes": n_brackets},
-        ]
+                for threshold_kwargs in thresholds:
+                    for multitask_handling in multitask_handlings:
+                        for target_to_int in target_to_ints:
+                            err_msg = f"{metric} - {target_nan_mask} - {threshold_kwargs}"
+
+                            if (multitask_handling is None) and (target_nan_mask == "ignore"):
+                                # Raise with incompatible options
+                                with self.assertRaises(ValueError):
+                                    MetricWrapper(
+                                        metric=metric,
+                                        threshold_kwargs=threshold_kwargs,
+                                        target_nan_mask=target_nan_mask,
+                                        multitask_handling=multitask_handling,
+                                        target_to_int=target_to_int,
+                                    )
+
+                            else:
+                                metric_wrapper = MetricWrapper(
+                                    metric=metric,
+                                    threshold_kwargs=threshold_kwargs,
+                                    target_nan_mask=target_nan_mask,
+                                    multitask_handling=multitask_handling,
+                                    target_to_int=target_to_int,
+                                )
+
+                                # Check that the metric can be saved and re-loaded without error
+                                torch.save(metric_wrapper, pickle_file)
+                                metric_wrapper2 = torch.load(pickle_file)
+                                self.assertTrue(metric_wrapper == metric_wrapper2, msg=err_msg)
+
+                                # Check that the metric only contains primitive types
+                                state = metric_wrapper.__getstate__()
+                                if state["threshold_kwargs"] is not None:
+                                    self.assertIsInstance(
+                                        state["threshold_kwargs"], dict, msg=err_msg
+                                    )
+                                if isinstance(metric, str):
+                                    self.assertIsInstance(state["metric"], str, msg=err_msg)
+
+
+    def test_update_compute_reset(self):
+        torch.manual_seed(42)
+        th = 0.7
+
+        # ---------- ACCURACY ----------
+        metric = MetricWrapper(
+                metric="accuracy", threshold_kwargs={"threshold": th, "operator": "greater"}, task="binary",
+            )
+        for batch_size in [1, 5, 25, 100]:
+            # Generate random predictions and targets, and compute the true accuracy
+            preds = torch.rand(100, dtype=torch.float32)
+            target = torch.rand(100, dtype=torch.float32)
+            preds_greater = preds > th
+            target_greater = target > th
+            true_accuracy = (preds_greater == target_greater).float().mean()
+
+            # Test the reset, update and compute
+            metric.reset()
+            for ii in range(0, 100, batch_size):
+                preds_batch = preds[ii : ii + batch_size]
+                target_batch = target_greater[ii : ii + batch_size]
+                metric.update(preds_batch, target_batch)
+
+            self.assertAlmostEqual(metric.compute(), true_accuracy, places=5, msg=f"Error for batch_size={batch_size}")
+        
+        # ---------- PEARSONR ----------
+        metric = MetricWrapper(
+                metric="pearsonr",
+            )
+        for batch_size in [1, 5, 25, 100]:
+            # Generate random predictions and targets, and compute the true pearsonr
+            preds = torch.rand(100, dtype=torch.float32)
+            target = torch.rand(100, dtype=torch.float32)
+            true_pearson = pearson_corrcoef(preds, target)
+
+            # Test the reset, update and compute with pearsonr
+            metric.reset()
+            for ii in range(0, 100, batch_size):
+                preds_batch = preds[ii : ii + batch_size]
+                target_batch = target[ii : ii + batch_size]
+                metric.update(preds_batch, target_batch)
 
-        for metric, kwargs, expected_score in zip(metrics, other_kwargs, expected_scores):
-            metric_wrapper = MetricWrapper(
-                metric=metric,
+            self.assertAlmostEqual(metric.compute().numpy(), true_pearson.numpy(), places=5, msg=f"Error for batch_size={batch_size}")
+
+
+        # ---------- PEARSONR with mean-per-label ----------
+
+        metric = MetricWrapper(
+                metric="pearsonr",
                 multitask_handling="mean-per-label",
-                squeeze_targets=True,
+            )
+        for batch_size in [1, 5, 25, 100]:
+            # Generate random predictions and targets, and compute the true pearsonr
+            preds = torch.rand(100, 10, dtype=torch.float32)
+            target = torch.rand(100, 10, dtype=torch.float32)
+            true_pearson = pearson_corrcoef(preds, target).mean().numpy()
+
+            # Test the pearson reset, update and compute with mean-per-label
+            metric.reset()
+            for ii in range(0, 100, batch_size):
+                preds_batch = preds[ii : ii + batch_size]
+                target_batch = target[ii : ii + batch_size]
+                metric.update(preds_batch, target_batch)
+
+            self.assertAlmostEqual(metric.compute().numpy(), true_pearson, places=5, msg=f"Error for batch_size={batch_size}")
+
+        # ---------- AUROC with mean-per-label ----------
+        metric = MetricWrapper(
+                metric="auroc",
                 target_to_int=True,
-                **kwargs,
+                multitask_handling="mean-per-label",
+                task="binary",
             )
-            score = metric_wrapper(preds, target)
-
-            assert score == expected_score
+        for batch_size in [1, 5, 25, 100]:
+            # Generate random predictions and targets, and compute the true auroc
+            preds = torch.rand(100, 10, dtype=torch.float32)
+            target = (0.5*preds + 0.5*torch.rand(100, 10, dtype=torch.float32)) > th
+            true_auroc = torch.stack([auroc(preds[:, ii], target[:, ii], task="binary") for ii in range(preds.shape[1])]).mean().numpy()
+
+            # Test the auroc reset, update and compute with mean-per-label
+            metric.reset()
+            for ii in range(0, 100, batch_size):
+                preds_batch = preds[ii : ii + batch_size]
+                target_batch = target[ii : ii + batch_size]
+                metric.update(preds_batch, target_batch)
+
+            self.assertAlmostEqual(metric.compute().numpy(), true_auroc, places=5, msg=f"Error for batch_size={batch_size}")
+        
 
 
 if __name__ == "__main__":
diff --git a/tests/test_multitask_datamodule.py b/tests/test_multitask_datamodule.py
index 3f7b09456..81a51459f 100644
--- a/tests/test_multitask_datamodule.py
+++ b/tests/test_multitask_datamodule.py
@@ -25,7 +25,7 @@
 TEMP_CACHE_DATA_PATH = "tests/temp_cache_0000"
 
 
-class Test_Multitask_DataModule(ut.TestCase):
+class test_Multitask_DataModule(ut.TestCase):
     def setUp(self):
         # Create a temporary directory
         self.tmp_test_dir = tempfile.mkdtemp()
diff --git a/tests/test_predictor.py b/tests/test_predictor.py
index 1ef69775f..3d3b8648c 100644
--- a/tests/test_predictor.py
+++ b/tests/test_predictor.py
@@ -29,7 +29,7 @@ def test_parse_loss_fun(self):
         preds = torch.rand(10, 5)
         target = (torch.rand(10, 5) > 0.5).to(preds.dtype)
         for this_loss in losses:
-            loss_fun = EvalOptions.parse_loss_fun(this_loss)
+            loss_name, loss_fun = EvalOptions.parse_loss_fun(this_loss)
             loss = loss_fun(preds, target)
 
 
diff --git a/tests/test_predictor_summaries.py b/tests/test_predictor_summaries.py
new file mode 100644
index 000000000..4e946d68b
--- /dev/null
+++ b/tests/test_predictor_summaries.py
@@ -0,0 +1,276 @@
+"""
+--------------------------------------------------------------------------------
+Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
+
+Use of this software is subject to the terms and conditions outlined in the LICENSE file.
+Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
+warranties of any kind.
+
+Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
+Refer to the LICENSE file for the full terms and conditions.
+--------------------------------------------------------------------------------
+"""
+
+
+"""
+Unit tests for the file graphium/trainer/predictor.py
+"""
+
+import torch
+from torch import nn
+from torchmetrics import MeanAbsoluteError, PearsonCorrCoef
+from copy import deepcopy
+import unittest as ut
+
+from graphium.trainer.predictor_summaries import MultiTaskSummary, STDMetric, GradientNormMetric
+
+class SimpleNN(nn.Module):
+# Define a simple neural network with 2 layers
+    def __init__(self, in_dim=10, out_dim=1):
+        super(SimpleNN, self).__init__()
+        torch.random.manual_seed(42)
+        # Define the first layer with 10 input features and 5 output features
+        self.layer1 = nn.Linear(in_dim, 5)
+        # Define the second layer with 5 input features and 1 output feature
+        self.layer2 = nn.Linear(5, out_dim)
+
+    def forward(self, x):
+        # Pass the input through the first layer
+        if x.ndim == 1:
+            x = x.unsqueeze(-1)
+        x = torch.relu(self.layer1(x))
+        # Pass the output of the first layer through the second layer
+        x = self.layer2(x)
+        return x
+    
+
+class SimpleDictNN(nn.Module):
+    def __init__(self, task_list, in_dim=10, out_dim=1):
+        super(SimpleDictNN, self).__init__()
+        torch.random.manual_seed(42)
+        self.dict_nn = nn.ModuleDict({task: SimpleNN(in_dim, out_dim) for task in task_list})
+
+    def forward(self, x):
+        return {task: self.dict_nn[task](x[task]) for task in self.dict_nn.keys()}
+
+
+def simple_nn_grad_step(model, inputs, targets):
+    # Initialize the optimizer and loss function
+    loss_fn = nn.MSELoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+    # Perform a gradient step
+    optimizer.zero_grad()
+    outputs = model(inputs)
+    if isinstance(outputs, dict):
+        loss = sum([loss_fn(outputs[task], targets[task]) for task in outputs.keys()])
+    else:
+        loss = loss_fn(outputs, targets)
+    loss.backward()
+    optimizer.step()
+    return model
+
+class test_TaskSummary(ut.TestCase):
+
+    def test_std_metric(self):
+            
+        # Generate random data
+        torch.random.manual_seed(42)
+        rand = torch.rand(100, 1)
+
+        # Compute expected values for STD
+        expected_std = torch.std(rand, correction=0)
+
+        # Compute std metric
+        std_metric = STDMetric()
+        std_metric.update(rand)
+        std_metric_val = std_metric.compute()
+        std_metric.reset()
+
+        self.assertAlmostEqual(std_metric_val.item(), expected_std.item(), places=5)
+
+        # Check multiple updates
+        std_metric.update(rand[:10])
+        std_metric.update(rand[10:25])
+        std_metric.update(rand[25:])
+        std_metric_val = std_metric.compute()
+        std_metric.reset()
+
+        self.assertAlmostEqual(std_metric_val.item(), expected_std.item(), places=5)
+
+        # Add some correction
+        expected_std = torch.std(rand, correction=1)
+        std_metric = STDMetric(correction=1)
+        std_metric.update(rand)
+        std_metric_val = std_metric.compute()
+        std_metric.reset()
+
+        self.assertAlmostEqual(std_metric_val.item(), expected_std.item(), places=5)
+
+        # Add some nans
+        rand[[3, 5, 11, 23, 42, 56, 78, 99]] = float('nan')
+        expected_std = torch.std(rand[~rand.isnan()], correction=0)
+
+        std_metric = STDMetric(nan_strategy='ignore', correction=0)
+        std_metric.update(rand)
+        std_metric_val = std_metric.compute()
+        std_metric.reset()
+
+        self.assertAlmostEqual(std_metric_val.item(), expected_std.item(), places=5)
+        
+    def test_gradient_norm_metric(self):
+
+        # Generate random data
+        torch.random.manual_seed(42)
+        LEN = 10000
+        inputs = torch.rand(LEN, 10)
+        targets = torch.rand(LEN, 1)
+
+        # Compute expected values for gradient norm
+        model = SimpleNN()
+        model = simple_nn_grad_step(model, inputs, targets)
+        expected_grad_norm = torch.norm(torch.stack([torch.norm(param.grad) for param in model.parameters()]))
+
+        # Compute gradient norm metric
+        model = SimpleNN()
+        model = simple_nn_grad_step(model, inputs, targets)
+        grad_norm_metric = GradientNormMetric()
+        grad_norm_metric.update(model)
+        grad_norm_metric_val = grad_norm_metric.compute()
+        grad_norm_metric.reset()
+
+        self.assertAlmostEqual(grad_norm_metric_val.item(), expected_grad_norm.item(), places=5)
+
+        # Compute gradient norm metric with many update steps
+        grad_norm_metric = GradientNormMetric()
+        model = SimpleNN()
+        model = simple_nn_grad_step(model, inputs[:50], targets[:50])
+        grad_norm_metric.update(model)
+        model = SimpleNN()
+        model = simple_nn_grad_step(model, inputs[50:400], targets[50:400])
+        grad_norm_metric.update(model)
+        model = SimpleNN()
+        model = simple_nn_grad_step(model, inputs[400:], targets[400:])
+        grad_norm_metric.update(model)
+
+        grad_norm_metric_val = grad_norm_metric.compute()
+        grad_norm_metric.reset()
+
+        self.assertAlmostEqual(grad_norm_metric_val.item(), expected_grad_norm.item(), places=1)
+
+    def assertDictTensorAlmostEqual(self, dict1, dict2, places=7):
+        dict1 = deepcopy(dict1)
+        dict1 = {key: dict1[key] for key in sorted(dict1.keys())}
+        dict2 = deepcopy(dict2)
+        dict2 = {key: dict2[key] for key in sorted(dict2.keys())}
+        for key in dict1.keys():
+            dict1[key] = round(dict1[key].item(), places)
+        for key in dict2.keys():
+            dict2[key] = round(dict2[key].item(), places)
+        self.assertDictEqual(dict1, dict2)
+
+
+    def test_multi_task_summary(self):
+
+        # Generate random data
+        torch.random.manual_seed(42)
+        targets = torch.rand(100, 3)
+        preds = torch.rand(100, 3) + 0.4 * targets
+        targets = {f"task{i+1}": targets[:, i] for i in range(targets.shape[1])}
+        preds = {f"task{i+1}": preds[:, i] for i in range(preds.shape[1])}
+
+        task_metrics = {
+            "task1": {'mae': MeanAbsoluteError(), 'pearson': PearsonCorrCoef()},
+            "task2": {'pearson': PearsonCorrCoef()},
+            "task3": {'mae': MeanAbsoluteError()}
+        }
+
+        expected_dict = {}
+        for task, metrics in task_metrics.items():
+            for metric_name, metric in metrics.items():
+                metric.update(preds[task], targets[task])
+                expected_val = metric.compute()
+                metric.reset()
+                expected_dict[f"{task}/{metric_name}/val"] = expected_val
+        
+
+        # Test the metrics on validation step
+        summary_val = MultiTaskSummary(task_metrics, step_name="val", compute_mean=False, compute_std=False)
+        summary_val.update(preds, targets)
+        summary_dict = summary_val.compute()
+        self.assertDictTensorAlmostEqual(summary_dict, expected_dict, places=5)
+
+        # Test the metric reset
+        summary_val.reset()
+        summary_val.update(preds, targets)
+        summary_dict = summary_val.compute()
+        self.assertDictTensorAlmostEqual(summary_dict, expected_dict, places=5)
+
+        # Test multiple batches
+        summary_val.reset()
+        preds1 = {key: preds[key][:10] for key in preds.keys()}
+        targets1 = {key: targets[key][:10] for key in targets.keys()}
+        preds2 = {key: preds[key][10:25] for key in preds.keys()}
+        targets2 = {key: targets[key][10:25] for key in targets.keys()}
+        preds3 = {key: preds[key][25:] for key in preds.keys()}
+        targets3 = {key: targets[key][25:] for key in targets.keys()}
+        
+        summary_val.update(preds1, targets1)
+        summary_val.update(preds2, targets2)
+        summary_val.update(preds3, targets3)
+        summary_dict = summary_val.compute()
+        self.assertDictTensorAlmostEqual(summary_dict, expected_dict, places=5)
+
+        # Test the mean and std computation
+        summary_val = MultiTaskSummary(task_metrics, step_name="val", compute_mean=True, compute_std=True)
+        summary_val.update(preds, targets)
+        summary_dict = summary_val.compute()
+        expected_dict_mean_std = {}
+        for task in task_metrics.keys():
+            expected_dict_mean_std[f"{task}/mean_preds/val"] = preds[task].mean()
+            expected_dict_mean_std[f"{task}/std_preds/val"] = preds[task].std(correction=0)
+            expected_dict_mean_std[f"{task}/mean_target/val"] = targets[task].mean()
+            expected_dict_mean_std[f"{task}/std_target/val"] = targets[task].std(correction=0)
+        expected_dict_mean_std.update(expected_dict)
+        self.assertDictTensorAlmostEqual(summary_dict, expected_dict_mean_std, places=5)
+
+        # Test the mean and std computation with multiple batches
+        summary_val.reset()
+        summary_val.update(preds1, targets1)
+        summary_val.update(preds2, targets2)
+        summary_val.update(preds3, targets3)
+        summary_dict = summary_val.compute()
+        self.assertDictTensorAlmostEqual(summary_dict, expected_dict_mean_std, places=5)
+
+        # Test the training step doesn't return anything when no metrics on training set are selected
+        summary_train = MultiTaskSummary(task_metrics, step_name="train", task_metrics_on_training_set=None, compute_mean=False, compute_std=False)
+        summary_train.update(preds, targets)
+        summary_train = summary_train.compute()
+        self.assertDictEqual(summary_train, {})
+
+        # Test the training step returns only the mae
+        task_metrics_on_training_set = {"task1": ["mae"], "task2": None, "task3": "mae"}
+        summary_train = MultiTaskSummary(task_metrics, step_name="train", task_metrics_on_training_set=task_metrics_on_training_set, compute_mean=False, compute_std=False)
+        summary_train.update(preds, targets)
+        summary_dict = summary_train.compute()
+        expected_dict_mae = {key: value for key, value in expected_dict.items() if "mae" in key}
+        expected_dict_mae = {key.replace("/val", "/train"): value for key, value in expected_dict_mae.items()}
+        self.assertDictTensorAlmostEqual(summary_dict, expected_dict_mae, places=5)
+
+        # Test the training step returns only the mae with multiple steps
+        summary_train = MultiTaskSummary(task_metrics, step_name="train", task_metrics_on_training_set=task_metrics_on_training_set, compute_mean=False, compute_std=False)
+        summary_train.update(preds1, targets1)
+        summary_train.update(preds2, targets2)
+        summary_train.update(preds3, targets3)
+        summary_dict = summary_train.compute()
+        self.assertDictTensorAlmostEqual(summary_dict, expected_dict_mae, places=5)
+
+        # Test grad_norm not available in "val" step
+        summary_val = MultiTaskSummary(task_metrics, step_name="val", compute_mean=False, compute_std=False)
+        summary_val.update(preds, targets)
+        summary_dict = summary_val.compute()
+        self.assertNotIn("grad_norm", summary_dict.keys())
+
+
+if __name__ == "__main__":
+    ut.main()
diff --git a/tests/test_training.py b/tests/test_training.py
index 1e007b6b7..789bbca3a 100644
--- a/tests/test_training.py
+++ b/tests/test_training.py
@@ -17,10 +17,25 @@
 import sys
 import subprocess
 import os
-from unittest.mock import patch
+import shutil
+import unittest as ut
 
 
-class TestCLITraining:
+import hydra
+from hydra.core.global_hydra import GlobalHydra
+from omegaconf import OmegaConf
+import os
+
+from graphium.config._loader import (
+    load_accelerator,
+    load_architecture,
+    load_datamodule,
+    load_metrics,
+    load_predictor,
+    load_trainer,
+)
+
+class test_CLITraining():
     @classmethod
     def setup_class(cls):
         print("Setting up the test class...")
@@ -98,8 +113,8 @@ def test_cpu_cli_training(self):
     @pytest.mark.ipu
     @pytest.mark.skip
     def test_ipu_cli_training(self):
-        with patch("poptorch.ipuHardwareIsAvailable", return_value=True):
-            with patch("lightning_graphcore.accelerator._IPU_AVAILABLE", new=True):
+        with ut.patch("poptorch.ipuHardwareIsAvailable", return_value=True):
+            with ut.patch("lightning_graphcore.accelerator._IPU_AVAILABLE", new=True):
                 import poptorch
 
                 assert poptorch.ipuHardwareIsAvailable()
@@ -107,3 +122,89 @@ def test_ipu_cli_training(self):
 
                 assert _IPU_AVAILABLE is True
                 self.call_cli_with_overrides("ipu", "16-true")
+
+
+
+def initialize_hydra(config_path, job_name="app"):
+    if GlobalHydra.instance().is_initialized():
+        GlobalHydra.instance().clear()
+    hydra.initialize(config_path=config_path, job_name=job_name)
+
+def compose_main_config(config_dir):
+    initialize_hydra(config_dir)
+    # Compose the main configuration
+    main_config = hydra.compose(config_name="main")
+    return main_config
+
+def compose_task_config(config_dir, task_name):
+    task_config_dir = os.path.join(config_dir, "tasks")
+    initialize_hydra(task_config_dir, job_name="compose_task")
+    # Compose the specific task configuration
+    task_config = hydra.compose(config_name=task_name)
+    return task_config
+
+class test_TrainToymix(ut.TestCase):
+    def test_train_toymix(self):
+        # Load the main configuration for toymix
+        CONFIG_DIR = "../expts/hydra-configs/"
+        cfg = compose_main_config(CONFIG_DIR)
+        cfg = OmegaConf.to_container(cfg, resolve=True)
+        cfg.pop("tasks")
+
+        # Adapt the configuration to reduce the time it takes to run the test, less samples, less epochs
+        cfg["constants"]["max_epochs"] = 4
+        cfg["trainer"]["trainer"]["check_val_every_n_epoch"] = 1
+        cfg["trainer"]["trainer"]["max_epochs"] = 4
+
+        cfg["datamodule"]["args"]["batch_size_training"] = 20
+        cfg["datamodule"]["args"]["batch_size_inference"] = 20
+        cfg["datamodule"]["args"]["task_specific_args"]["zinc"]["sample_size"] = 300
+        cfg["datamodule"]["args"]["task_specific_args"]["qm9"]["sample_size"] = 300
+        cfg["datamodule"]["args"]["task_specific_args"]["tox21"]["sample_size"] = 300
+
+        
+        # Initialize the accelerator
+        cfg, accelerator_type = load_accelerator(cfg)
+
+        # If the data_cache directory exists, delete it for the purpose of the test
+        data_cache = cfg["datamodule"]["args"]["processed_graph_data_path"]
+        if os.path.exists(data_cache):
+            shutil.rmtree(data_cache)
+
+        # Load and initialize the dataset
+        datamodule = load_datamodule(cfg, accelerator_type)
+
+        # Initialize the network
+        model_class, model_kwargs = load_architecture(
+            cfg,
+            in_dims=datamodule.in_dims,
+        )
+
+        datamodule.prepare_data()
+
+        metrics = load_metrics(cfg)
+
+        predictor = load_predictor(
+            cfg,
+            model_class,
+            model_kwargs,
+            metrics,
+            datamodule.get_task_levels(),
+            accelerator_type,
+            datamodule.featurization,
+            datamodule.task_norms,
+        )
+
+        metrics_on_progress_bar = predictor.get_metrics_on_progress_bar
+        trainer = load_trainer(cfg, accelerator_type, metrics_on_progress_bar=metrics_on_progress_bar)
+
+        predictor.set_max_nodes_edges_per_graph(datamodule, stages=["train", "val"])
+
+        # Run the model training
+        trainer.fit(model=predictor, datamodule=datamodule)
+        trainer.test(model=predictor, datamodule=datamodule)
+
+if __name__ == "__main__":
+    config_dir = "../expts/hydra-configs/"  # Path to your config directory
+
+    ut.main()
diff --git a/tests/test_utils.py b/tests/test_utils.py
index d14584798..537f35775 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -178,7 +178,7 @@ def test_safe_run(self):
             print("This is not an error")
 
 
-class TestTensorFp16ToFp32(ut.TestCase):
+class test_TensorFp16ToFp32(ut.TestCase):
     def test_tensor_fp16_to_fp32(self):
         # Create a tensor
         tensor = torch.randn(10, 10).half()